torch.cuda.comm.broadcast

torch.cuda.comm.broadcast(tensor, devices=None, *, out=None)

Broadcasts a tensor to specified GPU devices.

Parameters

• tensor (Tensor) – tensor to broadcast. Can be on CPU or GPU.

• devices (Iterable[torch.device, str or int], optional) – an iterable of GPU devices, among which to broadcast.

• out (Sequence[Tensor], optional, keyword-only) – the GPU tensors to store output results.

Note

Exactly one of devices and out must be specified.

Returns

• If devices is specified,

  a tuple containing copies of tensor, placed on devices.

• If out is specified,

  a tuple containing out tensors, each containing a copy of tensor.