Merge pull request huggingface#2229 from Promisery/reg_token

Initialize weights of reg_token for ViT
xwang233 · Jul 18, 2024 · 7160af4 · 7160af4
2 parents 34c9fee + 417cf7f
commit 7160af4
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/timm/models/vision_transformer.py b/timm/models/vision_transformer.py
@@ -590,6 +590,8 @@ def init_weights(self, mode: str = '') -> None:
             trunc_normal_(self.pos_embed, std=.02)
         if self.cls_token is not None:
             nn.init.normal_(self.cls_token, std=1e-6)
+        if self.reg_token is not None:
+            nn.init.normal_(self.reg_token, std=1e-6)
         named_apply(get_init_weights_vit(mode, head_bias), self)
 
     def _init_weights(self, m: nn.Module) -> None: