Merge pull request FlagOpen#1216 from Swgj/fix-issue#1215

fixed a bug in AbsReranker.py for mps device support
hanhainebula · Nov 13, 2024 · f3f9800 · f3f9800
2 parents c2a2625 + 04a9dc3
commit f3f9800
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/FlagEmbedding/abc/inference/AbsReranker.py b/FlagEmbedding/abc/inference/AbsReranker.py
@@ -96,7 +96,7 @@ def get_target_devices(devices: Union[str, int, List[str], List[int]]) -> List[s
             elif is_torch_npu_available():
                 return [f"npu:{i}" for i in range(torch.npu.device_count())]
             elif torch.backends.mps.is_available():
-                return [f"mps:{i}" for i in range(torch.mps.device_count())]
+                return ["mps"]
             else:
                 return ["cpu"]
         elif isinstance(devices, str):