@ptrblck Thanks a lot, you have no idea how may hours I spent on this. If you don’t mind me asking I have another question.
In the following scenario:
densenet.features.denseblock4.denselayer32(x)

I have the following error:
 self.densenet.features.denseblock4.denselayer32(x)
 *** RuntimeError: running_mean should contain 3 elements not 1888

Is that because I have to access each layer/part of the network via the forward function?