From 755b9dcb778deb5a82029d69ae3260579c6450f1 Mon Sep 17 00:00:00 2001 From: askender Date: Sun, 10 Apr 2022 00:59:27 +0800 Subject: [PATCH] Converting the list to ndarray for speedup (#4470) --- parlai/agents/rag/retrievers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/parlai/agents/rag/retrievers.py b/parlai/agents/rag/retrievers.py index d164c015525..f7d472612d1 100644 --- 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