[BugFix] Fix tests failing because of pytorch/pytorch#137602

cc mikaylagawarecki albanD

ghstack-source-id: 6fc7434
Pull Request resolved: #2558

Author: Vincent Moens

test/test_trainer.py

@@ -450,7 +450,7 @@ def make_storage():
             rb_trainer2.register(trainer2)
             if re_init:
                 trainer2._process_batch_hook(td.to_tensordict().zero_())
-            trainer2.load_from_file(file)
+            trainer2.load_from_file(file, weights_only=False)
             assert state_dict_has_been_called[0]
             assert load_state_dict_has_been_called[0]
             assert state_dict_has_been_called_td[0]

torchrl/trainers/trainers.py

@@ -296,12 +296,17 @@ def save_trainer(self, force_save: bool = False) -> None:
         if _save and self.save_trainer_file:
             self._save_trainer()
 
-    def load_from_file(self, file: Union[str, pathlib.Path]) -> Trainer:
+    def load_from_file(self, file: Union[str, pathlib.Path], **kwargs) -> Trainer:
+        """Loads a file and its state-dict in the trainer.
+
+        Keyword arguments are passed to the :func:`~torch.load` function.
+
+        """
         if _CKPT_BACKEND == "torchsnapshot":
             snapshot = Snapshot(path=file)
             snapshot.restore(app_state=self.app_state)
         elif _CKPT_BACKEND == "torch":
-            loaded_dict: OrderedDict = torch.load(file)
+            loaded_dict: OrderedDict = torch.load(file, **kwargs)
             self.load_state_dict(loaded_dict)
         return self