torch.cuda.comm.broadcast_coalesced

torch.cuda.comm.broadcast_coalesced(tensors, devices, buffer_size=10485760)

將一系列 Tensor 廣播到指定的 GPU。

小型的 Tensor 會先合併到一個緩衝區中，以減少同步的次數。

參數

• tensors (序列) – 要廣播的 Tensor。必須位於相同的裝置上，無論是 CPU 還是 GPU。

• devices (Iterable[torch.device, str or int]) – 要在其間廣播的 GPU 裝置的 Iterable。

• buffer_size (int) – 用於合併的緩衝區的最大大小

傳回

一個包含 tensor 副本的元組，放置在 devices 上。