Commit dc0238b8 authored by Adán Mora-Fallas's avatar Adán Mora-Fallas Committed by Francisco Massa
Browse files

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
parent ee570b66
......@@ -106,7 +106,7 @@ class ResNet(nn.Module):
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():
......
......@@ -82,7 +82,7 @@ class SqueezeNet(nn.Module):
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():
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment