From 0889a5b8907e8f0e14723a7942b8900573da3ee5 Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Tue, 23 Apr 2024 16:28:00 +0100 Subject: [PATCH] init --- test/test_transforms.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/test_transforms.py b/test/test_transforms.py index 36408bf4964..c9d2fb8c031 100644 --- 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)