Skip to content

GitLab

"next_docs/en/git@developer.sourcefind.cn:wangsen/mineru.git" did not exist on "4bb543939e1c4f1d9c4034e3c845dda082f962ae"
Unverified Commit dd9fb896 authored by Stella Biderman's avatar Stella Biderman Committed by GitHub
Browse files

Merge pull request #447 from nikhilpinnaparaju/issue-437 

Updated handling for device in lm_eval/models/gpt2.py
parents 602abceb 7376b0fd
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 3 deletions
lm_eval/models/gpt2.py
View file @ dd9fb896
...@@ -23,9 +23,8 @@ class HFLM(BaseLM): ...@@ -23,9 +23,8 @@ class HFLM(BaseLM):
assert isinstance(pretrained, str) assert isinstance(pretrained, str)
assert isinstance(batch_size, int) assert isinstance(batch_size, int)
if device: device_list = set(["cuda", "cpu"] + [f'cuda:{i}' for i in range(torch.cuda.device_count())])
if device not in ["cuda", "cpu"]: if device and device in device_list:
device = int(device)
self._device = torch.device(device) self._device = torch.device(device)
print(f"Using device '{device}'") print(f"Using device '{device}'")
else: else:
......
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