torch.cuda.comm.gather

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

從多個 GPU 裝置收集 tensors。

參數

• tensors (Iterable[Tensor]) – 要收集的 tensors 的可迭代物件。除了 dim 之外，所有維度的 Tensor 大小都必須匹配。

• dim (int, optional) – tensors 將被串接的維度。預設值：0。

• destination (torch.device, str, or int, optional) – 輸出裝置。可以是 CPU 或 CUDA。預設值：目前的 CUDA 裝置。

• out (Tensor, optional, keyword-only) – 用於儲存收集結果的 tensor。除了 dim 之外，其大小必須與 tensors 的大小匹配，其中大小必須等於 sum(tensor.size(dim) for tensor in tensors)。可以在 CPU 或 CUDA 上。

注意

當指定 out 時，不得指定 destination。

返回值

• 如果指定了 destination，

  一個位於 destination 裝置上的 tensor，它是沿著 dim 串連 tensors 的結果。

• 如果指定了 out，

  則 out tensor，現在包含沿著 dim 串連 tensors 的結果。