fix: visiontransformer

jina-ai · Nov 21, 2022 · 460622f · 460622f
1 parent 8a36b8a
commit 460622f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/server/clip_server/model/model.py b/server/clip_server/model/model.py
@@ -69,8 +69,8 @@ def __init__(self, layers: int, dtype: torch.dtype = torch.float32, **kwargs):
 
 
 class VisionTransformer(_VisionTransformer):
-    def __init__(self, dtype: torch.dtype = torch.float32, **kwargs):
-        super().__init__(**kwargs)
+    def __init__(self, image_size: int, patch_size: int, output_dim: int ,dtype: torch.dtype = torch.float32, **kwargs):
+        super().__init__(image_size, patch_size, output_dim=output_dim, **kwargs)
         self.transformer = Transformer(dtype=dtype, **kwargs)