issue#126: torch.load device issue.

This commit is contained in:
Kunpeng GUO
2023-02-25 07:09:07 +01:00
parent e19ee681ac
commit 85ad682530
+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: