"vscode:/vscode.git/clone" did not exist on "3f942acfe15de367931a63aa96c1931eb74799c0"
Commit 2bb7ce3b authored by baberabb's avatar baberabb
Browse files

add mps

parent 73b2149f
...@@ -99,7 +99,7 @@ class HFLM(LM): ...@@ -99,7 +99,7 @@ class HFLM(LM):
if not (parallelize or accelerator.num_processes > 1): if not (parallelize or accelerator.num_processes > 1):
# use user-passed device # use user-passed device
device_list = set( device_list = set(
["cuda", "cpu"] ["cuda", "cpu", "mps"]
+ [f"cuda:{i}" for i in range(torch.cuda.device_count())] + [f"cuda:{i}" for i in range(torch.cuda.device_count())]
) )
if device: if device:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment