[BugFix] Fix sampling in NonTensorSpec (#2172)

pytorch · May 26, 2024 · 78b4717 · 78b4717
1 parent a93063b
commit 78b4717
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/torchrl/data/tensor_specs.py b/torchrl/data/tensor_specs.py
@@ -1953,21 +1953,21 @@ def rand(self, shape=None):
         if shape is None:
             shape = ()
         return NonTensorData(
-            data=None, batch_size=(*shape, self.shape), device=self.device
+            data=None, batch_size=(*shape, *self.shape), device=self.device
         )
 
     def zero(self, shape=None):
         if shape is None:
             shape = ()
         return NonTensorData(
-            data=None, batch_size=(*shape, self.shape), device=self.device
+            data=None, batch_size=(*shape, *self.shape), device=self.device
         )
 
     def one(self, shape=None):
         if shape is None:
             shape = ()
         return NonTensorData(
-            data=None, batch_size=(*shape, self.shape), device=self.device
+            data=None, batch_size=(*shape, *self.shape), device=self.device
         )
 
     def is_in(self, val: torch.Tensor) -> bool: