Converting the list to ndarray for speedup (#4470)

facebookresearch · Apr 9, 2022 · 755b9dc · 755b9dc
1 parent 1d70b0e
commit 755b9dc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/parlai/agents/rag/retrievers.py b/parlai/agents/rag/retrievers.py
@@ -656,7 +656,7 @@ def index_retrieve(
             query.cpu().detach().to(torch.float32).numpy(), n_docs
         )
         ids, np_vectors = zip(*top_docs_and_scores)
-        vectors = torch.tensor(np_vectors).to(query)
+        vectors = torch.tensor(np.array(np_vectors)).to(query)
         if isinstance(self.indexer, DenseHNSWFlatIndexer):
             vectors = vectors[:, :, :-1]
         # recompute exact FAISS scores