[BugFix] Fix serialization of stacks of Tensorclasses

ghstack-source-id: 8e47f46
Pull Request resolved: #1236

(cherry picked from commit 635c9c0)

Author: Vincent Moens

tensordict/_lazy.py

@@ -381,7 +381,7 @@ def from_dict(
         stack_dim_name=None,
         stack_dim=0,
     ):
-        return LazyStackedTensorDict(
+        return cls._new_lazy_unsafe(
             *(
                 TensorDict.from_dict(
                     input_dict[str(i)],

tensordict/_reductions.py

@@ -91,8 +91,12 @@ def from_metadata(metadata=metadata, prefix=None):
         _ = metadata.pop("size", None)
 
         d = {
-            key: NonTensorData(data, batch_size=batch_size)
-            for (key, (data, batch_size)) in non_tensor.items()
+            key: NonTensorData(
+                data,
+                batch_size=batch_size,
+                device=torch.device(device) if device is not None else None,
+            )
+            for (key, (data, batch_size, device)) in non_tensor.items()
         }
         for key, (dtype, local_shape, start, stop, pad) in leaves.items():
             dtype = _STRDTYPE2DTYPE[dtype]

tensordict/base.py

@@ -5055,6 +5055,7 @@ def assign(
                     metadata_dict["non_tensors"][key] = (
                         value.data,
                         list(value.batch_size),
+                        str(value.device) if value.device is not None else None,
                     )
                 return
             elif _is_tensor_collection(cls):

test/test_tensorclass.py

@@ -233,6 +233,12 @@ class MyDataClass:
     MyTensorClass_autocast = MyTensorClass_nocast = MyTensorClass = None
 
 
+@tensorclass
+class TCStrings:
+    a: str
+    b: str
+
+
 class TestTensorClass:
     def test_get_default(self):
         @tensorclass
@@ -1250,6 +1256,21 @@ def test_pickle(self):
         assert isinstance(data2, MyData)
         assert data2.z == data.z
 
+    @pytest.mark.parametrize("consolidate", [False, True])
+    def test_pickle_consolidate(self, consolidate):
+        with set_capture_non_tensor_stack(False):
+
+            tc = TCStrings(a="a", b="b")
+
+            tcstack = TensorDict(tc=torch.stack([tc, tc.clone()]))
+            if consolidate:
+                tcstack = tcstack.consolidate()
+            assert isinstance(tcstack["tc"], TCStrings)
+            loaded = pickle.loads(pickle.dumps(tcstack))
+            assert isinstance(loaded["tc"], TCStrings)
+            assert loaded["tc"].a == tcstack["tc"].a
+            assert loaded["tc"].b == tcstack["tc"].b
+
     def test_post_init(self):
         @tensorclass
         class MyDataPostInit: