reference/generated/tensordict.TensorDictBase.from_remote_init

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.from_remote_init

classmethod TensorDictBase.from_remote_init(src: int, group: 'ProcessGroup' | None = None, device: device | None = None, use_broadcast: bool = False) → Self

Receives a tensordict that init_remote() sends, and returns it.

The result is consolidated (see consolidate()): receive the next batches into it with recv(consolidated=True), which allocates no new memory. A tensorclass arrives as the same tensorclass, if this process can import its class. The metadata is unpickled: use it only with ranks you trust. See init_remote() for an example, and Distributed tensordicts.

Parameters:

• src (int) – the rank of the sender.

• group (torch.distributed.ProcessGroup, optional) – the process group. Defaults to None, the default process group.

• device (device, str or int, optional) – the device of the result. Point-to-point, it is also the device that the metadata goes through, as in torch.distributed.recv_object_list(). Defaults to None, which is the CPU point-to-point, and the device of the sender’s storage with use_broadcast=True.

• use_broadcast (bool, optional) – if True, receives with broadcast(). It must match the sender’s value. Defaults to False.

Returns:

A new consolidated tensordict, of the class that was sent.

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