diff --git a/torchvision/models/densenet.py b/torchvision/models/densenet.py index 0a9376d1b7a..43f8e7c3eba 100644 --- a/torchvision/models/densenet.py +++ b/torchvision/models/densenet.py @@ -112,7 +112,7 @@ def __init__(self, growth_rate=32, block_config=(6, 12, 24, 16), def forward(self, x): features = self.features(x) out = F.relu(features, inplace=True) - out = F.avg_pool2d(out, kernel_size=7, stride=1).view(features.size(0), -1) + out = F.adaptive_avg_pool2d(out, (1, 1)).view(features.size(0), -1) out = self.classifier(out) return out