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) – 用於儲存收集結果的張量。除了 dim 外，其大小必須與 tensors 的大小匹配，在 dim 維度上，大小必須等於 sum(tensor.size(dim) for tensor in tensors)。可以在 CPU 或 CUDA 上。

注意

指定 out 時，不得指定 destination。

返回值

• 如果指定了 destination，

  一個位於 destination 裝置上的張量，它是將 tensors 沿 dim 連線的結果。

• 如果指定了 out，

  out 張量，現在包含將 tensors 沿 dim 連線的結果。