Merge pull request #45 from sbintuitions/fix/reranking_tensor

[Fix] rerankingでのformat不具合
sbintuitions · Jul 30, 2024 · 161d40a · 161d40a
2 parents fe590c1 + 624ceeb
commit 161d40a
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/src/jmteb/evaluators/reranking/evaluator.py b/src/jmteb/evaluators/reranking/evaluator.py
@@ -135,9 +135,11 @@ def _compute_metrics(
             reranked_docs_list = []
             for i, item in enumerate(query_dataset):
                 query_embedding = to_tensor(query_embeddings[i], device=device)
-                doc_embedding = to_tensor(
-                    torch.stack([doc_embeddings[doc_indices[retrieved_doc]] for retrieved_doc in item.retrieved_docs]),
-                    device=device,
+                doc_embedding = torch.stack(
+                    [
+                        Tensor(doc_embeddings[doc_indices[retrieved_doc]]).to(device=device)
+                        for retrieved_doc in item.retrieved_docs
+                    ]
                 )
                 similarity = dist_func(query_embedding, doc_embedding)