faiss index on gpu

webis-de · Jan 1, 2025 · 38519e6 · 38519e6
1 parent 086c2d8
commit 38519e6
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/lightning_ir/retrieve/faiss/faiss_searcher.py b/lightning_ir/retrieve/faiss/faiss_searcher.py
@@ -14,12 +14,19 @@
 
 class FaissSearcher(Searcher):
     def __init__(
-        self, index_dir: Path | str, search_config: FaissSearchConfig, module: BiEncoderModule, use_gpu: bool = False
+        self,
+        index_dir: Path | str,
+        search_config: FaissSearchConfig,
+        module: BiEncoderModule,
+        use_gpu: bool = False,
+        use_gpu_index: bool = False,
     ) -> None:
         import faiss
 
         self.search_config: FaissSearchConfig
         self.index = faiss.read_index(str(Path(index_dir) / "index.faiss"))
+        if use_gpu and use_gpu_index and hasattr(faiss, "index_cpu_to_all_gpus"):
+            self.index = faiss.index_cpu_to_all_gpus(self.index)
         ivf_index = None
         try:
             ivf_index = faiss.extract_index_ivf(self.index)