[BugFix] Faster and safer non-tensor stack

ghstack-source-id: f6c61cd14954a87a6de333eda2cd14e5f822a4e9 Pull Request resolved: #1232
pytorch · Feb 24, 2025 · e23ce5c · e23ce5c
1 parent d25bd54
commit e23ce5c
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tensordict/tensorclass.py b/tensordict/tensorclass.py
@@ -3227,6 +3227,8 @@ def _stack_non_tensor(cls, list_of_non_tensor, dim=0, raise_if_non_unique=False)
         ids = set()
         firstdata = NO_DEFAULT
         return_stack = not capture_non_tensor_stack()
+        if return_stack:
+            return NonTensorStack(*list_of_non_tensor, stack_dim=dim)
         for data in list_of_non_tensor:
             if not isinstance(data, NonTensorData):
                 if raise_if_non_unique: