Update davit.py

Author: fffffgggg54

timm/models/davit.py

@@ -518,6 +518,7 @@ def __init__(
         self.head = nn.Sequential(OrderedDict([
             ('global_pool', SelectAdaptivePool2d(pool_type=global_pool)),
             ('norm', nn.Identity() if head_norm_first else norm_layer(self.num_features)),
+            ('flatten', nn.Flatten(1) if global_pool else nn.Identity()),
             ('drop', nn.Dropout(self.drop_rate)),
             ('fc', nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity())]))
 
@@ -545,6 +546,7 @@ def get_classifier(self):
     def reset_classifier(self, num_classes, global_pool=None):
         if global_pool is not None:
             self.head.global_pool = SelectAdaptivePool2d(pool_type=global_pool)
+            self.head.flatten = nn.Flatten(1) if global_pool else nn.Identity()
         self.head.fc = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
 
     def forward_features(self, x):
@@ -559,7 +561,7 @@ def forward_features(self, x):
     def forward_head(self, x, pre_logits: bool = False):
         x = self.head.global_pool(x)
         x = self.head.norm(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
-        x = x.squeeze()
+        x = self.head.flatten(x)
         x = self.head.drop(x)
         return x if pre_logits else self.head.fc(x)