amend

pytorch · Feb 15, 2024 · a6dc5ee · a6dc5ee
1 parent 70c1523
commit a6dc5ee
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/test/test_env.py b/test/test_env.py
@@ -2523,7 +2523,7 @@ def _step(self, tensordict):
             )
 
     torch.manual_seed(0)
-    policy = Actor(torch.nn.Linear(3, 3))
+    policy = Actor(torch.nn.Linear(3, 3, device=device))
     env0 = DifferentiableEnv(device=device)
     seed = env0.set_seed(0)
     env1 = DifferentiableEnv(device=device)