Changing to AdaptiveAvgPool2d on SqueezeNet and ResNet (#643)

* AdaptiveAvgpool2d integrated

The Average Pool Layer on SqueezeNet and ResNet was hardcoded, this was changed, now any the input size is accepted.

* error fixed

* linting fixed

Author: AdanMora

torchvision/models/resnet.py

@@ -106,7 +106,7 @@ def __init__(self, block, layers, num_classes=1000):
         self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
         self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
         self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
-        self.avgpool = nn.AvgPool2d(7, stride=1)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
         self.fc = nn.Linear(512 * block.expansion, num_classes)
 
         for m in self.modules():

torchvision/models/squeezenet.py

@@ -82,7 +82,7 @@ def __init__(self, version=1.0, num_classes=1000):
             nn.Dropout(p=0.5),
             final_conv,
             nn.ReLU(inplace=True),
-            nn.AvgPool2d(13, stride=1)
+            nn.AdaptiveAvgPool2d((1, 1))
         )
 
         for m in self.modules():