Fix dtype for searcher pids

stanford-futuredata · Nov 17, 2023 · b97e216 · b97e216
1 parent 03c8fc8
commit b97e216
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/colbert/search/index_storage.py b/colbert/search/index_storage.py
@@ -90,7 +90,7 @@ def rank(self, config, Q, filter_fn=None, pids=None):
                 pids, centroid_scores = self.retrieve(config, Q)
             else:
                 centroid_scores = None
-                pids = torch.tensor(pids, dtype=pids.dtype, device=Q.device)
+                pids = torch.tensor(pids, dtype=torch.int32, device=Q.device)
 
             if filter_fn is not None:
                 filtered_pids = filter_fn(pids)