Merge remote-tracking branch 'origin/main' into master-patched

xwang233 · Jul 18, 2024 · 6c3d339 · 6c3d339
2 parents a5b6f97 + 7160af4
commit 6c3d339
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: