[Feature] Composite.batch_size

ghstack-source-id: 621884a
Pull Request resolved: #2597

Author: Vincent Moens

torchrl/data/tensor_specs.py

@@ -4257,6 +4257,14 @@ def __new__(cls, *args, **kwargs):
         cls._locked = False
         return super().__new__(cls)
 
+    @property
+    def batch_size(self):
+        return self._shape
+
+    @batch_size.setter
+    def batch_size(self, value: torch.Size):
+        self._shape = value
+
     @property
     def shape(self):
         return self._shape