Update

[ghstack-poisoned]
pytorch · Feb 25, 2025 · 1706004 · 1706004
1 parent 2aa9170
commit 1706004
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/tensordict/_reductions.py b/tensordict/_reductions.py
@@ -91,8 +91,8 @@ 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=device)
+            for (key, (data, batch_size, device)) in non_tensor.items()
         }
         for key, (dtype, local_shape, start, stop, pad) in leaves.items():
             dtype = _STRDTYPE2DTYPE[dtype]

diff --git a/tensordict/base.py b/tensordict/base.py
@@ -5056,6 +5056,7 @@ def assign(
                     metadata_dict["non_tensors"][key] = (
                         value.data,
                         list(value.batch_size),
+                        value.device,
                     )
                 return
             elif _is_tensor_collection(cls):