fix: sample size

jina-ai · Oct 15, 2021 · 1dc475f · 1dc475f
1 parent f47e366
commit 1dc475f
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/finetuner/tuner/evaluation.py b/finetuner/tuner/evaluation.py
@@ -17,7 +17,7 @@ def extract_catalog(docs):
 
 
 def prepare_eval_docs(docs, catalog, limit=10, sample_size=100, seed=42):
-    sampled_docs = docs.sample(sample_size, seed)
+    sampled_docs = docs.sample(min(sample_size, len(docs)), seed)
     to_be_scored_docs = DocumentArray()
     for doc in sampled_docs:
         d = Document(

diff --git a/finetuner/tuner/pytorch/__init__.py b/finetuner/tuner/pytorch/__init__.py
@@ -194,7 +194,7 @@ def get_embeddings(self, data: DocumentArrayLike):
         blobs = data.blobs
         embeddings = self.embed_model(torch.Tensor(blobs).to(self.device))
         for doc, embed in zip(data, embeddings):
-            doc.embedding = np.array(embed)
+            doc.embedding = np.array(embed.cpu())
 
     def save(self, *args, **kwargs):
         torch.save(self.embed_model.state_dict(), *args, **kwargs)