fix npu support method in single card

EleutherAI · May 16, 2024 · cbbef02 · cbbef02
1 parent 0f4fe94
commit cbbef02
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/lm_eval/models/huggingface.py b/lm_eval/models/huggingface.py
@@ -164,9 +164,9 @@ def __init__(
                 # use user-passed device
                 device_list = set(
                     ["cuda", "cpu", "npu"]
-                    + [f"cuda:{i}" for i in range(device_counts)]
-                    + [f"npu:{i}" for i in range(device_counts)]
+                    + [f"cuda:{i}" for i in range(torch.cuda.device_count())]
                     + ["mps", "mps:0"]
+                    + ["npu", "npu:0"]
                 )
                 if device and device in device_list:
                     self._device = torch.device(device)