fix reset head in hgnet

Author: seefun

timm/models/hgnet.py

@@ -470,8 +470,8 @@ def reset_classifier(self, num_classes, global_pool='avg'):
                 class_expand=self.class_expand,
                 use_lab=self.use_lab)
         else:
-            if self.global_pool == 'avg':
-                self.head = SelectAdaptivePool2d(pool_type=self.global_pool, flatten=True)
+            if global_pool == 'avg':
+                self.head = SelectAdaptivePool2d(pool_type=global_pool, flatten=True)
             else:
                 self.head = nn.Identity()
 
@@ -480,7 +480,7 @@ def forward_features(self, x):
         return self.stages(x)
 
     def forward_head(self, x, pre_logits: bool = False):
-        return self.head(x, pre_logits=pre_logits)
+        return self.head(x, pre_logits=pre_logits) if pre_logits else self.head(x)
 
     def forward(self, x):
         x = self.forward_features(x)