amend

pytorch · vmoens · Feb 5, 2024 · Feb 1, 2024 · Feb 1, 2024 · Feb 2, 2024
commit 0e80f1ba234c6abe222005601325db32f2b27879
diff --git a/test/mocking_classes.py b/test/mocking_classes.py
@@ -1072,7 +1072,7 @@ def _step(
         tensordict: TensorDictBase,
     ) -> TensorDictBase:
         action = tensordict.get(self.action_key)
-        self.count += action.to(torch.int).to(self.device)
+        self.count += action.to(dtype=torch.int, device=self.device)
         tensordict = TensorDict(
             source={
                 "observation": self.count.clone(),

diff --git a/test/test_collector.py b/test/test_collector.py
@@ -1690,8 +1690,8 @@ def test_reset_heterogeneous_envs(
         env_device = torch.device("cpu")  # explicit mapping
     elif env_device is not None and env_device.type == "cuda" and policy_device is None:
         policy_device = torch.device("cpu")
-    env1 = lambda: TransformedEnv(CountingEnv(), StepCounter(2))
-    env2 = lambda: TransformedEnv(CountingEnv(), StepCounter(3))
+    env1 = lambda: TransformedEnv(CountingEnv(device="cpu"), StepCounter(2))
+    env2 = lambda: TransformedEnv(CountingEnv(device="cpu"), StepCounter(3))
     if parallel:
         cls = ParallelEnv
     else: