fix

pytorch · Jan 3, 2024 · db9cac4 · db9cac4
1 parent 7cf7496
commit db9cac4
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/tensordict/_lazy.py b/tensordict/_lazy.py
@@ -53,6 +53,7 @@
     infer_size_impl,
     is_tensorclass,
     KeyedJaggedTensor,
+    lazy_legacy,
     lock_blocked,
     NestedKey,
 )
@@ -785,6 +786,13 @@ def _get_str(
                         self._batch_size
                         + out.batch_size[(len(self._batch_size) + incr) :]
                     )
+                elif not lazy_legacy():
+                    # it must be a TensorDict
+                    incr = 0 if not self._is_vmapped else 1
+                    out._batch_size = (
+                        self._batch_size
+                        + out.batch_size[(len(self._batch_size) + incr) :]
+                    )
                 else:
                     # then it's a tensorclass
                     out._tensordict.hook_out = self.hook_out