torch.cuda.comm.gather

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

将多个 GPU 设备上的张量汇集在一起。

参数

• tensors (Iterable[Tensor]) – 一个可迭代的张量集合，用于汇集。除 dim 维度外，所有其他维度的张量大小必须匹配。

• dim (int, optional) – 将张量连接到的维度。默认为 0。

• destination (torch.device, str, or int, optional) – 输出设备。可以是 CPU 或 CUDA。默认为当前的 CUDA 设备。

• out (Tensor, optional, keyword-only) – 用于存储汇集结果的张量。其大小必须与 tensors 匹配，除了 dim 维度，该维度的大小必须等于 sum(tensor.size(dim) for tensor in tensors)。可以是 CPU 或 CUDA。

注意

当指定 out 时，不能指定 destination。

返回

• 如果指定了 destination，

  则返回一个位于 destination 设备上的张量，它是将 tensors 沿 dim 维度连接的结果。

• 如果指定了 out，

  则返回 out 张量，该张量现在包含将 tensors 沿 dim 维度连接的结果。