Update

[ghstack-poisoned]
pytorch · Jan 7, 2025 · fef37d6 · fef37d6
2 parents bce850d + e07daee
commit fef37d6
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tensordict/tensorclass.py b/tensordict/tensorclass.py
@@ -938,6 +938,9 @@ def wrapper(
             batch_size = kwargs.pop("batch_size", torch.Size(()))
         if isinstance(batch_size, int):
             batch_size = (batch_size,)
+        elif batch_size is None:
+            batch_size = torch.Size(())
+
         if "names" in required_params:
             names = None
         else: