Update

[ghstack-poisoned]
pytorch · Aug 11, 2024 · f56b8e0 · f56b8e0
1 parent 65ca9b8
commit f56b8e0
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/test/mocking_classes.py b/test/mocking_classes.py
@@ -1039,7 +1039,8 @@ def _step(
     ) -> TensorDictBase:
         action = tensordict.get(self.action_key)
         self.count += action.to(
-            dtype=torch.int, device=self.full_observation_spec["observation"].device
+            dtype=torch.int,
+            device=self.action_spec.device if self.device is None else self.device,
         )
         tensordict = TensorDict(
             source={