diff --git a/pretrainedmodels/models/senet.py b/pretrainedmodels/models/senet.py index 20bd122f..bc669b91 100644 --- a/pretrainedmodels/models/senet.py +++ b/pretrainedmodels/models/senet.py @@ -320,7 +320,7 @@ def __init__(self, block, layers, groups, reduction, dropout_p=0.2, downsample_kernel_size=downsample_kernel_size, downsample_padding=downsample_padding ) - self.avg_pool = nn.AvgPool2d(7, stride=1) + self.avg_pool = nn.AdaptiveAvgPool2d(1) self.dropout = nn.Dropout(dropout_p) if dropout_p is not None else None self.last_linear = nn.Linear(512 * block.expansion, num_classes)