Use weights_only for load (huggingface#1933)

Co-authored-by: Quentin Gallouédec <45557362+qgallouedec@users.noreply.github.com>

Author: kit1980

examples/scripts/ddpo.py

@@ -93,7 +93,7 @@ def __init__(self, *, dtype, model_id, model_filename):
             cached_path = hf_hub_download(model_id, model_filename)
         except EntryNotFoundError:
             cached_path = os.path.join(model_id, model_filename)
-        state_dict = torch.load(cached_path, map_location=torch.device("cpu"))
+        state_dict = torch.load(cached_path, map_location=torch.device("cpu"), weights_only=True)
         self.mlp.load_state_dict(state_dict)
         self.dtype = dtype
         self.eval()

tests/test_peft_models.py

@@ -138,7 +138,7 @@ def test_save_pretrained_peft(self):
             assert os.path.exists(f"{tmp_dir}/adapter_config.json"), f"{tmp_dir}/adapter_config.json does not exist"
             # check also for `pytorch_model.bin` and make sure it only contains `v_head` weights
             assert os.path.exists(f"{tmp_dir}/pytorch_model.bin"), f"{tmp_dir}/pytorch_model.bin does not exist"
-            maybe_v_head = torch.load(f"{tmp_dir}/pytorch_model.bin")
+            maybe_v_head = torch.load(f"{tmp_dir}/pytorch_model.bin", weights_only=True)
             # check that only keys that starts with `v_head` are in the dict
             assert all(
                 k.startswith("v_head") for k in maybe_v_head.keys()

trl/models/auxiliary_modules.py

@@ -60,7 +60,7 @@ def __init__(self, *, dtype, model_id, model_filename):
             cached_path = hf_hub_download(model_id, model_filename)
         except EntryNotFoundError:
             cached_path = os.path.join(model_id, model_filename)
-        state_dict = torch.load(cached_path, map_location=torch.device("cpu"))
+        state_dict = torch.load(cached_path, map_location=torch.device("cpu"), weights_only=True)
         self.mlp.load_state_dict(state_dict)
         self.dtype = dtype
         self.eval()