[BugFix] Fix error for incongruent devices in load_state_dict (#529)

pytorch · Sep 14, 2023 · a721766 · a721766
1 parent 0006b91
commit a721766
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/tensordict/tensordict.py b/tensordict/tensordict.py
@@ -815,11 +815,8 @@ def load_state_dict(
         state_dict = copy(state_dict)
         self.batch_size = state_dict.pop("__batch_size")
         device = state_dict.pop("__device", None)
-        if device is not None:
-            if device != self.device:
-                raise RuntimeError(
-                    "Loading data from another device is not yet supproted."
-                )
+        if device is not None and self.device is not None and device != self.device:
+            raise RuntimeError("Loading data from another device is not yet supported.")
 
         for key, item in state_dict.items():
             if isinstance(item, dict):