[BugFix] Faster and safer non-tensor stack

ghstack-source-id: f6c61cd
Pull Request resolved: #1232

(cherry picked from commit e23ce5c)

Author: Vincent Moens

tensordict/tensorclass.py

@@ -3219,7 +3219,9 @@ def _stack_non_tensor(cls, list_of_non_tensor, dim=0, raise_if_non_unique=False)
 
         ids = set()
         firstdata = NO_DEFAULT
-        return_stack = False
+        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: