Update

[ghstack-poisoned]

Author: Vincent Moens

tensordict/_lazy.py

@@ -1538,7 +1538,9 @@ def densify(self, *, layout: torch.layout = torch.strided):
                             list_of_entries, layout=layout
                         )
                     else:
-                        raise NotImplementedError
+                        raise NotImplementedError(
+                            f"stack_dim is {self.stack_dim} but not 0. Densify canot be done."
+                        )
             else:
                 tensor = self._get_str(key, None)
                 if tensor is not None:

test/test_tensorclass.py

@@ -15,6 +15,7 @@
 import re
 import sys
 import weakref
+from dataclasses import field
 from multiprocessing import Pool
 from pathlib import Path
 from tempfile import TemporaryDirectory
@@ -744,6 +745,17 @@ class MyClass2:
         assert (a != c.clone().zero_()).any()
         assert (c != a.clone().zero_()).any()
 
+    def test_field(self):
+        class Cls(TensorClass):
+            a: torch.Tensor
+            b: str
+            c: dict = field(default_factory=dict)
+
+        obj = Cls(a=torch.arange(3), b="abc", batch_size=[3])
+        assert obj[0].a == obj[1].a - 1
+        assert obj[0].b == obj[1].b
+        assert obj[0].c is obj[1].c
+
     def test_from_dataclass(self):
         assert is_tensorclass(MyTensorClass_autocast)
         assert MyTensorClass_nocast is not MyDataClass