update

gau-nernst · Oct 29, 2023 · 711b18e · 711b18e
1 parent 13dbac7
commit 711b18e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/vision_toolbox/backbones/vit.py b/vision_toolbox/backbones/vit.py
@@ -169,17 +169,18 @@ def from_config(variant: str, img_size: int, *, weights: str | None = None) -> V
 
         return m
 
-    # https://github.com/google-research/vision_transformer
     @torch.no_grad()
     def load_jax_ckpt(self, ckpt: str, big_vision: bool = False, prefix: str = "") -> None:
         if big_vision:
+            # https://github.com/google-research/big_vision
             gcs_bucket = "big_vision"
             mha_norm = "LayerNorm_0"
             mha = "MultiHeadDotProductAttention_0"
             mlp_norm = "LayerNorm_1"
             mlp = "MlpBlock_0"
 
         else:
+            # https://github.com/google-research/vision_transformer
             gcs_bucket = "vit_models"
             mha_norm = "LayerNorm_0"
             mha = "MultiHeadDotProductAttention_1"