reference/generated/tensordict.TensorDictBase.all_reduce

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.all_reduce

TensorDictBase.all_reduce(op=None, *, group: 'torch.distributed.ProcessGroup' | None = None, async_op: bool = False) → None

All-reduces the tensordict across all ranks in-place.

Each leaf tensor is reduced individually via dist.all_reduce. After the call, every rank holds the reduced values.

Parameters:

op (dist.ReduceOp, optional) – The reduce operation (e.g. ReduceOp.SUM). Defaults to ReduceOp.SUM.

Keyword Arguments:

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

• async_op (bool) – if True, returns a list of futures. Defaults to False.

Returns:

None, or a list of futures if async_op=True.

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