self.device in huggingface.py line 210 treated as torch.device but mi…

…ght be a string (EleutherAI#1172) * self.device in huggingface.py line 210 In huggingface.py line 210, self.device is str and does not have a "type" attribute * Update huggingface.py This handles both the case where `self.device` is a `torch.device` and a string * Update huggingface.py --------- Co-authored-by: Hailey Schoelkopf <[email protected]>
mozilla-ai · Dec 19, 2023 · 78545d4 · 78545d4
1 parent 3f0a361
commit 78545d4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/lm_eval/models/huggingface.py b/lm_eval/models/huggingface.py
@@ -170,7 +170,7 @@ def __init__(
                         f"Using `accelerate launch` or `parallelize=True`, device '{device}' will be overridden when placing model."
                     )
                 # TODO: include in warning that `load_in_8bit` etc. affect this too
-                self._device = device
+                self._device = torch.device(device)
 
             # TODO: update this to be less of a hack once subfolder is fixed in HF
             revision = revision + ("/" + subfolder if subfolder is not None else "")
@@ -207,7 +207,7 @@ def __init__(
         self.model.eval()
         self.model.tie_weights()
 
-        if (gpus >= 1 or self.device.type == "mps") and isinstance(pretrained, str):
+        if (gpus >= 1 or str(self.device) == "mps") and isinstance(pretrained, str):
             if not (parallelize or autogptq or ("device_map" in kwargs)):
                 # place model onto device requested manually,
                 # if not using HF Accelerate or device_map