fix: test

michaelfeil · Jun 22, 2024 · 3b71a9c · 3b71a9c
1 parent 76901c1
commit 3b71a9c
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/libs/infinity_emb/tests/unit_test/transformer/vision/test_torch_vision.py b/libs/infinity_emb/tests/unit_test/transformer/vision/test_torch_vision.py
@@ -1,3 +1,4 @@
+import numpy as np
 import requests  # type: ignore
 import torch
 from PIL import Image  # type: ignore
@@ -22,8 +23,11 @@ def test_clip_like_model():
         image,
     ]
     embeddings = model.encode_post(model.encode_core(model.encode_pre(inputs)))
-    assert len(embeddings) == len(inputs)
 
+    assert isinstance(embeddings, list)
+    assert isinstance(embeddings[0], np.ndarray)
+    assert len(embeddings) == len(inputs)
+    embeddings = torch.tensor(embeddings)
     model = CLIPModel.from_pretrained(model_name)
     processor = CLIPProcessor.from_pretrained(model_name)