Solved by my3bikaht in post #4 
              
              
                (this option ain’t pretty, it’s more like a hack). For example here’s a forward method of resnet class: 
    def _forward_impl(self, x):
        #out = {}
        x = self.conv1(x)
        #out['conv1'] = x.detach()
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
       …