TensorDictBase.pin_memory#

TensorDictBase.pin_memory(num_threads: int | None = None, inplace: bool = False) → Self#

Calls pin_memory() on the stored tensors.

Parameters:

num_threads (int, optional) – if provided, the number of threads to use to call pin_memory on the leaves, in a ThreadPoolExecutor. Defaults to None, which, like 0, calls pin_memory() on the calling thread.
inplace (bool, optional) – if True, the tensordict is modified in-place. When every node is an unlocked TensorDict or LazyStackedTensorDict (or a tensorclass over one) and no leaf is a MemoryMappedTensor, in shared memory, sparse, jagged or an UnbatchedTensor, each leaf is replaced by its pinned copy, and the tensordict and every nested node stay the same objects; keys that hold the same tensor share its pinned copy, and a leaf that is a view of another leaf gets its own. Any other tree, such as a locked one or one that holds a memory-mapped tensordict, has the pinned copies written back into its leaves, as before, so they stay in pageable memory. Defaults to False.

Every tensordict class and every tensorclass has this method. For the other methods, see Methods by task.

Docs