PQFS into Index trainer (facebookresearch#3941)

Summary: Pull Request resolved: facebookresearch#3941 Support PQFS with index trainer Reviewed By: kuarora Differential Revision: D64259953 fbshipit-source-id: fd7ed90aed2ff6b6351460dcf7b61058c59cd25b
KamathForAIX · Oct 12, 2024 · 6441b56 · 6441b56
1 parent 847cde8
commit 6441b56
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/benchs/bench_fw/index.py b/benchs/bench_fw/index.py
@@ -73,7 +73,10 @@ def set_index_param_dict(index, param_dict, assert_same=False):
     def set_index_param(index, name, val, assert_same=False):
         index = faiss.downcast_index(index)
         val = int(val)
-        if isinstance(index, faiss.IndexPreTransform):
+        if (
+            isinstance(index, faiss.IndexPreTransform)
+            or isinstance(index, faiss.IndexIDMap)
+        ):
             Index.set_index_param(index.index, name, val)
             return
         elif name == "snap":