[Performance] Faster to

ghstack-source-id: 3ff1db59f081b75f24c34c5239f88b5c5de8dbe4 Pull Request resolved: #1073
pytorch · Nov 5, 2024 · b1339d1 · b1339d1
1 parent f12d31d
commit b1339d1
Showing 1 changed file with 13 additions and 1 deletion.
diff --git a/tensordict/base.py b/tensordict/base.py
@@ -10873,7 +10873,19 @@ def to_pinmem(tensor, _to=to):
                         to_pinmem, propagate_lock=True, **apply_kwargs
                     )
                 else:
-                    result = result._fast_apply(to, propagate_lock=True, **apply_kwargs)
+                    # result = result._fast_apply(to, propagate_lock=True, **apply_kwargs)
+                    keys, tensors = self._items_list(True, True)
+                    tensors = [to(t) for t in tensors]
+                    items = dict(zip(keys, tensors))
+                    result = self._fast_apply(
+                        lambda name, val: items.get(name, val),
+                        named=True,
+                        nested_keys=True,
+                        is_leaf=_NESTED_TENSORS_AS_LISTS,
+                        propagate_lock=True,
+                        **apply_kwargs,
+                    )
+
         if batch_size is not None:
             result.batch_size = batch_size
         if (