init

pytorch · Apr 23, 2024 · 0889a5b · 0889a5b
1 parent 7dd0128
commit 0889a5b
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/test/test_transforms.py b/test/test_transforms.py
@@ -10583,6 +10583,7 @@ def test_multistep_transform(self):
 
         outs_2 = []
         td = env.reset().contiguous()
+        assert "reward" not in td
         for _ in range(1):
             rollout = env.rollout(
                 250, auto_reset=False, tensordict=td, break_when_any_done=False
@@ -10626,7 +10627,7 @@ def test_multistep_transform(self):
             ).contiguous()
             assert "reward" not in rollout.keys()
             out = t._inv_call(rollout)
-            td = rollout[..., -1]["next"]
+            td = rollout[..., -1]["next"].exclude("reward")
             if out is not None:
                 outs_3.append(out)