Fix Binary reshaping

Author: Luca Carminati

test/test_specs.py

@@ -3913,6 +3913,7 @@ def test_expand(self):
         disc = Categorical(shape=(-1, 1, 2), n=4)
         moneh = MultiOneHot(shape=(-1, 1, 2, 7), nvec=[3, 4])
         mdisc = MultiCategorical(shape=(-1, 1, 2, 2), nvec=[3, 4])
+        binary = Binary(shape=(-1, 1, 2))
 
         spec = Composite(
             unb=unb,
@@ -3922,6 +3923,7 @@ def test_expand(self):
             disc=disc,
             moneh=moneh,
             mdisc=mdisc,
+            binary=binary,
             shape=(-1, 1, 2),
         )
         assert spec.shape == (-1, 1, 2)

torchrl/data/tensor_specs.py

@@ -4449,15 +4449,15 @@ def expand(self, *shape):
                 f"shape of the {self.__class__.__name__} spec in expand()."
             )
         return self.__class__(
-            n=self.shape[-1] if len(self.shape) > 0 else None,
+            n=shape[-1] if len(shape) > 0 else None,
             shape=shape,
             device=self.device,
             dtype=self.dtype,
         )
 
     def _reshape(self, shape):
         return self.__class__(
-            n=self.shape[-1] if len(self.shape) > 0 else None,
+            n=shape[-1] if len(shape) > 0 else None,
             shape=shape,
             device=self.device,
             dtype=self.dtype,
@@ -4470,7 +4470,7 @@ def _unflatten(self, dim, sizes):
             .shape
         )
         return self.__class__(
-            n=self.shape[-1] if len(self.shape) > 0 else None,
+            n=shape[-1] if len(shape) > 0 else None,
             shape=shape,
             device=self.device,
             dtype=self.dtype,
@@ -4481,7 +4481,7 @@ def squeeze(self, dim=None):
         if shape is None:
             return self
         return self.__class__(
-            n=self.shape[-1] if len(self.shape) > 0 else None,
+            n=shape[-1] if len(shape) > 0 else None,
             shape=shape,
             device=self.device,
             dtype=self.dtype,
@@ -4490,7 +4490,7 @@ def squeeze(self, dim=None):
     def unsqueeze(self, dim: int):
         shape = _unsqueezed_shape(self.shape, dim)
         return self.__class__(
-            n=self.shape[-1] if len(self.shape) > 0 else None,
+            n=shape[-1] if len(shape) > 0 else None,
             shape=shape,
             device=self.device,
             dtype=self.dtype,
@@ -4510,7 +4510,7 @@ def unbind(self, dim: int = 0):
         shape = tuple(s for i, s in enumerate(self.shape) if i != dim)
         return tuple(
             self.__class__(
-                n=self.shape[-1] if len(self.shape) > 0 else None,
+                n=shape[-1] if len(shape) > 0 else None,
                 shape=shape,
                 device=self.device,
                 dtype=self.dtype,