fix_freeze_without_cls_token_vit (#1693)

2023-07-14 09:43:19 +02:00 · 2023-07-14 09:43:19 +02:00 · db395d35b1
parent 465b6bdeec
commit db395d35b1
1 changed files with 2 additions and 1 deletions
--- a/mmpretrain/models/backbones/vision_transformer.py
+++ b/mmpretrain/models/backbones/vision_transformer.py
@ -436,7 +436,8 @@ class VisionTransformer(BaseBackbone):
        for param in self.pre_norm.parameters():
            param.requires_grad = False
        # freeze cls_token
-        self.cls_token.requires_grad = False
+        if self.cls_token:
+            self.cls_token.requires_grad = False
        # freeze layers
        for i in range(1, self.frozen_stages + 1):
            m = self.layers[i - 1]