torch.cuda.comm.gather

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

从多个GPU设备收集张量。

Parameters

• 张量 (可迭代[张量]) – 要收集的张量的可迭代对象。 除dim维度外的所有张量大小必须匹配。

• dim (int, 可选) – 张量将沿此维度连接。默认值：0。

• 目标 (torch.device, str, 或 int, 可选) – 输出设备。 可以是 CPU 或 CUDA。默认值：当前 CUDA 设备。

• out (Tensor, 可选, 仅关键字) – 用于存储收集结果的张量。 其大小必须与tensors的大小匹配，除了dim维度， 在该维度上的大小必须等于sum(tensor.size(dim) for tensor in tensors)。 可以在CPU或CUDA上。

注意

destination 在指定 out 时不能被指定。

Returns

• 如果指定了 destination，

  位于 destination 设备上的张量，即沿着 dim 连接 tensors 的结果。

• 如果指定了 out，

  out 张量，现在包含沿着 dim 连接 tensors 的结果。