Initialize weights of reg_token for ViT

Author: Promisery

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: