amend

pytorch · Apr 23, 2024 · 5f9828a · 5f9828a
1 parent 0889a5b
commit 5f9828a
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/torchrl/data/tensor_specs.py b/torchrl/data/tensor_specs.py
@@ -967,14 +967,14 @@ def zero(self, shape=None) -> TensorDictBase:
             dim = self.dim + len(shape)
         else:
             dim = self.dim
-        return torch.stack([spec.zero(shape) for spec in self._specs], dim)
+        return LazyStackedTensorDict.maybe_dense_stack([spec.zero(shape) for spec in self._specs], dim)
 
     def rand(self, shape=None) -> TensorDictBase:
         if shape is not None:
             dim = self.dim + len(shape)
         else:
             dim = self.dim
-        return torch.stack([spec.rand(shape) for spec in self._specs], dim)
+        return LazyStackedTensorDict.maybe_dense_stack([spec.rand(shape) for spec in self._specs], dim)
 
     def to(self, dest: Union[torch.dtype, DEVICE_TYPING]) -> T:
         if dest is None:
@@ -4229,7 +4229,7 @@ def project(self, val: TensorDictBase) -> TensorDictBase:
                 vals.append(spec.project(subval))
             else:
                 vals.append(subval)
-        res = torch.stack(vals, dim=self.dim)
+        res = LazyStackedTensorDict.maybe_dense_stack(vals, dim=self.dim)
         if not isinstance(val, LazyStackedTensorDict):
             res = res.to_tensordict()
         return res