[Refactor] Use td.transpose in multi-step transform (#2288)

pytorch · Jul 10, 2024 · 8e43ac8 · 8e43ac8
1 parent dcd332d
commit 8e43ac8
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/test/test_transforms.py b/test/test_transforms.py
@@ -10706,6 +10706,7 @@ def test_multistep_transform(self):
             rollout = env.rollout(
                 2, auto_reset=False, tensordict=td, break_when_any_done=False
             ).contiguous()
+            assert rollout.shape[:-1] == env.batch_size
             assert "reward" not in rollout.keys()
             out = t._inv_call(rollout)
             td = rollout[..., -1]["next"].exclude("reward")

diff --git a/torchrl/data/postprocs/postprocs.py b/torchrl/data/postprocs/postprocs.py
@@ -233,10 +233,7 @@ def _multi_step_func(
             try:
                 # let's try to reshape the tensordict
                 tensordict.batch_size = done.shape
-                tensordict = tensordict.apply(
-                    lambda x: x.transpose(ndim - 1, tensordict.ndim - 1),
-                    batch_size=done.transpose(ndim - 1, tensordict.ndim - 1).shape,
-                )
+                tensordict = tensordict.transpose(ndim - 1, tensordict.ndim - 1)
                 done = tensordict.get(("next", done_key))
             except Exception as err:
                 raise RuntimeError(