add mps

2bb7ce3b · baberabb · 73b2149f · 2bb7ce3b
Commit 2bb7ce3b authored Jul 15, 2023 by baberabb
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

lm_eval/models/huggingface.py lm_eval/models/huggingface.py +1 -1

No files found.
--- a/lm_eval/models/huggingface.py
+++ b/lm_eval/models/huggingface.py
@@ -99,7 +99,7 @@ class HFLM(LM):
        if not (parallelize or accelerator.num_processes > 1):
            # use user-passed device
            device_list = set(
-                ["cuda", "cpu"]
+                ["cuda", "cpu", "mps"]
                + [f"cuda:{i}" for i in range(torch.cuda.device_count())]
            )
            if device: