Merge pull request #134 from gabinguo/main

issue#126: torch.load device issue.
This commit is contained in:
Sourab Mangrulkar
2023-02-25 11:43:43 +05:30
committed by GitHub
+2 -1
View File
@@ -155,7 +155,8 @@ class PeftModel(PushToHubMixin, torch.nn.Module):
f"Please check that the file {WEIGHTS_NAME} is present at {model_id}."
)
adapters_weights = torch.load(filename)
adapters_weights = torch.load(
filename, map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu"))
# load the weights into the model
model = set_peft_model_state_dict(model, adapters_weights)
if getattr(model, "hf_device_map", None) is not None: