Update

[ghstack-poisoned]
pytorch · Jan 7, 2025 · 5dfc92a · 5dfc92a
1 parent 2140527
commit 5dfc92a
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tensordict/tensorclass.py b/tensordict/tensorclass.py
@@ -936,6 +936,9 @@ def wrapper(
             batch_size = torch.Size(())
         else:
             batch_size = kwargs.pop("batch_size", torch.Size(()))
+        if batch_size is None:
+            batch_size = torch.Size(())
+
         if "names" in required_params:
             names = None
         else: