reference/generated/tensordict.TensorDictBase.pin_memory

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.pin_memory

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

Returns a tensordict whose tensor leaves are in pinned (page-locked) memory.

Calls pin_memory() on every tensor leaf. Copies from pinned memory to an accelerator are faster, and can run asynchronously with to(device, non_blocking=True). Pinning needs an accelerator, such as CUDA. A PersistentTensorDict or a TensorDictStore raises a RuntimeError: call to_tensordict() first. See Devices, dtypes, and memory.

Parameters:

• num_threads (int, optional) – the number of threads that pin the leaves, in a ThreadPoolExecutor. Defaults to None, which, like 0, pins them on the calling thread.

• inplace (bool, optional) – if True, the tensordict and its nested tensordicts keep their identity, each leaf is replaced by its pinned copy, and the call returns self. A locked tensordict, one that holds a node other than a TensorDict, a lazy stack or a tensorclass, or one with memory-mapped, shared, sparse, jagged or unbatched leaves gets the pinned copies written back into its existing leaves, which stay in pageable memory. Defaults to False.

Examples

The example needs an accelerator, so it is not run when the documentation is tested.

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(obs=torch.zeros(1000, 3), batch_size=[1000])
>>> pinned = td.pin_memory()  
>>> pinned["obs"].is_pinned()  
True
>>> td_cuda = pinned.to("cuda", non_blocking=True)  

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