Update davit.py

huggingface · Dec 11, 2023 · 2597ce2 · 2597ce2
1 parent c82598f
commit 2597ce2
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)