Set weights_only in torch.load (#36991)
This commit is contained in:
@@ -554,7 +554,7 @@ class PeftIntegrationTester(unittest.TestCase, PeftTesterMixin):
|
||||
|
||||
state_dict_path = hf_hub_download(peft_model_id, "adapter_model.bin")
|
||||
|
||||
dummy_state_dict = torch.load(state_dict_path)
|
||||
dummy_state_dict = torch.load(state_dict_path, weights_only=True)
|
||||
|
||||
model.load_adapter(adapter_state_dict=dummy_state_dict, peft_config=peft_config)
|
||||
with self.assertRaises(ValueError):
|
||||
@@ -579,7 +579,7 @@ class PeftIntegrationTester(unittest.TestCase, PeftTesterMixin):
|
||||
|
||||
peft_config = LoraConfig()
|
||||
state_dict_path = hf_hub_download(peft_model_id, "adapter_model.bin")
|
||||
dummy_state_dict = torch.load(state_dict_path)
|
||||
dummy_state_dict = torch.load(state_dict_path, weights_only=True)
|
||||
|
||||
# this should always work
|
||||
model.load_adapter(
|
||||
@@ -647,7 +647,7 @@ class PeftIntegrationTester(unittest.TestCase, PeftTesterMixin):
|
||||
|
||||
peft_config = LoraConfig()
|
||||
state_dict_path = hf_hub_download(peft_model_id, "adapter_model.bin")
|
||||
dummy_state_dict = torch.load(state_dict_path)
|
||||
dummy_state_dict = torch.load(state_dict_path, weights_only=True)
|
||||
|
||||
# add unexpected key
|
||||
dummy_state_dict["foobar"] = next(iter(dummy_state_dict.values()))
|
||||
@@ -674,7 +674,7 @@ class PeftIntegrationTester(unittest.TestCase, PeftTesterMixin):
|
||||
|
||||
peft_config = LoraConfig()
|
||||
state_dict_path = hf_hub_download(peft_model_id, "adapter_model.bin")
|
||||
dummy_state_dict = torch.load(state_dict_path)
|
||||
dummy_state_dict = torch.load(state_dict_path, weights_only=True)
|
||||
|
||||
# remove a key so that we have missing keys
|
||||
key = next(iter(dummy_state_dict.keys()))
|
||||
|
||||
Reference in New Issue
Block a user