torch.cuda.comm.gather#

torch.cuda.comm.gather(tensors,dim=0,destination=None,*,out=None)[source]#

Gathers tensors from multiple GPU devices.

Parameters

• tensors (Iterable[Tensor]) – an iterable of tensors to gather.Tensor sizes in all dimensions other thandim have to match.

• dim (int,optional) – a dimension along which the tensors will beconcatenated. Default:0.

• destination (torch.device,str, orint,optional) – the output device.Can be CPU or CUDA. Default: the current CUDA device.

• out (Tensor,optional,keyword-only) – the tensor to store gather result.Its sizes must match those oftensors, except fordim,where the size must equalsum(tensor.size(dim)fortensorintensors).Can be on CPU or CUDA.

Note

destination must not be specified whenout is specified.

Returns

• Ifdestination is specified,

  a tensor located ondestination device, that is a result ofconcatenatingtensors alongdim.

• Ifout is specified,

  theout tensor, now containing results of concatenatingtensors alongdim.