torch.cuda.comm.broadcast

torch.cuda.comm.broadcast(tensor, devices=None, *, out=None)

將一個 tensor 廣播到指定的 GPU 裝置。

引數

• tensor (Tensor) – 要廣播的 tensor。可以是 CPU 或 GPU 上的。

• devices (Iterable[torch.device, str 或 int], optional) – 一個 GPU 裝置的可迭代物件，將在這些裝置之間進行廣播。

• out (Sequence[Tensor], optional, keyword-only) – 用於儲存輸出結果的 GPU tensor。

注意

必須且僅能指定 devices 和 out 中的一個。

返回

• 如果指定了 devices，

  則返回一個包含 tensor 副本的元組，這些副本放置在 devices 上。

• 如果指定了 out，

  則返回一個包含 out tensor 的元組，每個 tensor 都包含 tensor 的一個副本。