[Chore] set weights_only as False during loading for upcoming PyTorch…

… 2.6
ai4co · Jan 20, 2025 · ea4483f · ea4483f
1 parent e378ddc
commit ea4483f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/rl4co/models/rl/reinforce/reinforce.py b/rl4co/models/rl/reinforce/reinforce.py
@@ -201,9 +201,9 @@ def load_from_checkpoint(
             loaded.setup()
             loaded.post_setup_hook()
             # load baseline state dict
-            state_dict = torch.load(checkpoint_path, map_location=map_location)[
-                "state_dict"
-            ]
+            state_dict = torch.load(
+                checkpoint_path, map_location=map_location, weights_only=False
+            )["state_dict"]
             # get only baseline parameters
             state_dict = {k: v for k, v in state_dict.items() if "baseline" in k}
             state_dict = {k.replace("baseline.", "", 1): v for k, v in state_dict.items()}

diff --git a/rl4co/models/zoo/polynet/model.py b/rl4co/models/zoo/polynet/model.py
@@ -91,7 +91,7 @@ def __init__(
             logging.info(
                 f"Trying to load weights from baseline model {base_model_checkpoint_path}"
             )
-            checkpoint = torch.load(base_model_checkpoint_path)
+            checkpoint = torch.load(base_model_checkpoint_path, weights_only=False)
             state_dict = checkpoint["state_dict"]
             state_dict = {k.replace("policy.", "", 1): v for k, v in state_dict.items()}
             policy.load_state_dict(state_dict, strict=False)