fix npu support in single card

EleutherAI · May 16, 2024 · db5c319 · db5c319
1 parent 0f4fe94
commit db5c319
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/lm_eval/models/huggingface.py b/lm_eval/models/huggingface.py
@@ -163,10 +163,10 @@ def __init__(
             if not (parallelize or accelerator.num_processes > 1):
                 # 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)]
+                    ["cuda", "cpu"]
+                    + [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)