Hey @amohammadi,
You said that you have patched the dense net in some branch. Do you mind to open a MR for it?
Thanks