fix deit

gau-nernst · Oct 29, 2023 · c5f1889 · c5f1889
1 parent cad7d5b
commit c5f1889
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/vision_toolbox/backbones/deit.py b/vision_toolbox/backbones/deit.py
@@ -28,8 +28,8 @@ def __init__(
     ) -> None:
         # fmt: off
         super().__init__(
-            d_model, depth, n_heads, patch_size, img_size, True, bias, mlp_ratio,
-            dropout, layer_scale_init, stochastic_depth, norm_eps
+            d_model, depth, n_heads, patch_size, img_size, True, "cls_token", bias,
+            mlp_ratio, dropout, layer_scale_init, stochastic_depth, norm_eps,
         )
         # fmt: on
         self.dist_token = nn.Parameter(torch.zeros(1, 1, d_model))
@@ -133,7 +133,7 @@ def __init__(
     ):
         # fmt: off
         super().__init__(
-            d_model, depth, n_heads, patch_size, img_size, cls_token, bias,
+            d_model, depth, n_heads, patch_size, img_size, cls_token, "cls_token", bias,
             mlp_ratio, dropout, layer_scale_init, stochastic_depth, norm_eps,
         )
         # fmt: on