Solved by Kari_M.J in post #4 
              
              
                The problem was that I was missing nn.Flatten() before the fully connected layer. Changing the model head to 
        self.model_Head = nn.Sequential(
                                    nn.Conv2d(512, 256, kernel_size=1, stride=1, padding=0,bias=False),
                                    self.ResN…