all_gather#

composer.utils.dist.all_gather(tensor)[source]#

Collects a Tensor from each rank.

Parameters

tensor (Tensor) โ€“ Tensor from each rank to be gathered.

Returns

Sequence[Tensor] โ€“ A sequence of tensors indexed by rank.