[Feature] Composite.batch_size

ghstack-source-id: 621884a559a71e80a4be36c7ba984fd08be47952 Pull Request resolved: #2597
pytorch · Nov 24, 2024 · 2e82cab · 2e82cab
1 parent 097d8ad
commit 2e82cab
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/torchrl/data/tensor_specs.py b/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