Merge branch 'davit_head_update' into ml-decoder

huggingface · Dec 10, 2023 · 08b7837 · 08b7837
2 parents f59c1f1 + 318369a
commit 08b7837
Showing 1 changed file with 1 addition and 5 deletions.
diff --git a/timm/models/davit.py b/timm/models/davit.py
@@ -570,11 +570,7 @@ def forward_features(self, x):
         return x
 
     def forward_head(self, x, pre_logits: bool = False):
-        x = self.head.global_pool(x)
-        x = self.head.norm(x)
-        x = self.head.flatten(x)
-        x = self.head.drop(x)
-        return x if pre_logits else self.head.fc(x)
+        return self.head(x, pre_logits=True) if pre_logits else self.head(x)
 
     def forward(self, x):
         x = self.forward_features(x)