fix: Fix batch embedding test

qdrant · Nov 5, 2024 · e008a46 · e008a46
1 parent 2c4d320
commit e008a46
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_multi_task_models.py b/tests/test_multi_task_models.py
@@ -453,7 +453,7 @@ def test_batch_embedding():
             embedding_arrays = [e.embedding for e in embeddings]
             embeddings = np.stack(embedding_arrays)
 
-            assert np.allclose(embeddings, expected_vectors, atol=task_tolerance)
+            assert np.allclose(embeddings[:docs], expected_vectors, atol=task_tolerance)
 
         if is_ci:
             shutil.rmtree(model.model._model_dir)