Update vision_transformer.py

rmst · web-flow · commit 2eea97d64bbe · 2023-11-09T20:49:04.000-05:00
diff --git a/timm/models/vision_transformer.py b/timm/models/vision_transformer.py
@@ -682,9 +682,9 @@ def forward_head(self, x, pre_logits: bool = False):
         return x if pre_logits else self.head(x)
 
     def forward(self, x):
-        x = self.forward_features(x)
-        x = self.forward_head(x)
-        return x
+        tokens = self.forward_features(x)
+        cls_token = self.forward_head(tokens)
+        return cls_token, tokens
 
 
 def init_weights_vit_timm(module: nn.Module, name: str = ''):
