reference/generated/tensordict.TensorDictBase.share_memory_

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.share_memory

TensorDictBase.share_memory_() → Self

Moves the leaves to shared memory, locks the tensordict, and returns self.

Processes that receive the tensordict, for example through torch.multiprocessing, then read and write the same tensors. The lock forbids adding, removing or replacing entries, but in-place writes work (see Locking and temporary changes). On a tensordict whose device is CUDA, the call does nothing: CUDA tensors are shared already, and the tensordict is not locked. After unlock_(), is_shared() returns False, although the tensors stay in shared memory.

Raises:

RuntimeError – if the tensordict is memory-mapped.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(a=torch.zeros(3), batch_size=[3])
>>> td.share_memory_().is_shared(), td.is_locked, td["a"].is_shared()
(True, True, True)

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