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

Author: fedebotu

Diff for: 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 for: 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)