From 297d6b3a357aade0496cb29f7e6e5834e8a445d1 Mon Sep 17 00:00:00 2001 From: Vlad Lialin Date: Tue, 11 May 2021 23:42:30 -0400 Subject: [PATCH] simpllify faiss index save Fixes #2350 --- src/datasets/search.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/src/datasets/search.py b/src/datasets/search.py index f9b54e89b35..08b77df48ce 100644 --- a/src/datasets/search.py +++ b/src/datasets/search.py @@ -319,18 +319,11 @@ def save(self, file: Union[str, PurePath]): """Serialize the FaissIndex on disk""" import faiss # noqa: F811 - if ( - hasattr(self.faiss_index, "device") - and self.faiss_index.device is not None - and self.faiss_index.device > -1 - ) or ( - hasattr(self.faiss_index, "getDevice") - and self.faiss_index.getDevice() is not None - and self.faiss_index.getDevice() > -1 - ): + if self.device is not None and self.device > -1: index = faiss.index_gpu_to_cpu(self.faiss_index) else: index = self.faiss_index + faiss.write_index(index, str(file)) @classmethod