Prevent torch FutureWarning torch.load with weights_only=False wh…

…en loading checkpoints

skrl/agents/torch/base.py

@@ -6,6 +6,7 @@
 import os
 import gym
 import gymnasium
+from packaging import version
 
 import numpy as np
 import torch
@@ -374,7 +375,10 @@ def load(self, path: str) -> None:
         :param path: Path to load the model from
         :type path: str
         """
-        modules = torch.load(path, map_location=self.device)
+        if version.parse(torch.__version__) >= version.parse("1.13"):
+            modules = torch.load(path, map_location=self.device, weights_only=False)  # prevent torch:FutureWarning
+        else:
+            modules = torch.load(path, map_location=self.device)
         if type(modules) is dict:
             for name, data in modules.items():
                 module = self.checkpoint_modules.get(name, None)

skrl/models/torch/base.py

@@ -3,6 +3,7 @@
 import collections
 import gym
 import gymnasium
+from packaging import version
 
 import numpy as np
 import torch
@@ -488,7 +489,11 @@ def load(self, path: str) -> None:
             >>> model = Model(observation_space, action_space, device="cuda:1")
             >>> model.load("model.pt")
         """
-        self.load_state_dict(torch.load(path, map_location=self.device))
+        if version.parse(torch.__version__) >= version.parse("1.13"):
+            state_dict = torch.load(path, map_location=self.device, weights_only=False)  # prevent torch:FutureWarning
+        else:
+            state_dict = torch.load(path, map_location=self.device)
+        self.load_state_dict(state_dict)
         self.eval()
 
     def migrate(self,

skrl/models/torch/tabular.py

@@ -1,5 +1,7 @@
 from typing import Any, Mapping, Optional, Sequence, Tuple, Union
 
+from packaging import version
+
 import torch
 
 from skrl.models.torch import Model
@@ -196,7 +198,10 @@ def load(self, path: str) -> None:
             >>> model = Model(observation_space, action_space, device="cuda:1")
             >>> model.load("model.pt")
         """
-        tensors = torch.load(path)
+        if version.parse(torch.__version__) >= version.parse("1.13"):
+            tensors = torch.load(path, weights_only=False)  # prevent torch:FutureWarning
+        else:
+            tensors = torch.load(path)
         for name, tensor in tensors.items():
             if hasattr(self, name) and isinstance(getattr(self, name), torch.Tensor):
                 _tensor = getattr(self, name)

skrl/multi_agents/torch/base.py

@@ -6,6 +6,7 @@
 import os
 import gym
 import gymnasium
+from packaging import version
 
 import numpy as np
 import torch
@@ -393,7 +394,10 @@ def load(self, path: str) -> None:
         :param path: Path to load the model from
         :type path: str
         """
-        modules = torch.load(path, map_location=self.device)
+        if version.parse(torch.__version__) >= version.parse("1.13"):
+            modules = torch.load(path, map_location=self.device, weights_only=False)  # prevent torch:FutureWarning
+        else:
+            modules = torch.load(path, map_location=self.device)
         if type(modules) is dict:
             for uid in self.possible_agents:
                 if uid not in modules: