theguijo:

out = F.avg_pool2d(out, 4)


This is the problem. Original resnet uses adaptive_avg_pool2d(out, (1, 1)) in its last layer before connecting it to linear layer. So, if you change this line to
out = F.adaptive_avg_pool2d(out, (1, 1)),
your model will work just fine. This layer forces output to be in shape of output_size=(1, 1).
Here is a post about how it works:

  
    
    
    What is AdaptiveAvgPool2d? 
  
  
    The AdaptiveAvgPool2d layers confuse me a lot. 
Is there any math formula explaning it?
  


Here is the link to source code of this line:

  
      github.com
  
  
    pytorch/vision/blob/e6d3f8c515ee1a4bbda4c17594b00a4e5cb58e07/torchvision/models/resnet.py#L157 1

self.bn1 = norm_layer(self.inplanes)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = self._make_layer(block, 64, layers[0])
self.layer2 = self._make_layer(block, 128, layers[1], stride=2,
                               dilate=replace_stride_with_dilation[0])
self.layer3 = self._make_layer(block, 256, layers[2], stride=2,
                               dilate=replace_stride_with_dilation[1])
self.layer4 = self._make_layer(block, 512, layers[3], stride=2,
                               dilate=replace_stride_with_dilation[2])
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.fc = nn.Linear(512 * block.expansion, num_classes)

for m in self.modules():
    if isinstance(m, nn.Conv2d):
        nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
    elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
        nn.init.constant_(m.weight, 1)
        nn.init.constant_(m.bias, 0)

# Zero-initialize the last BN in each residual branch,



  
  
    
    
  
  


bests