torch.cuda.comm.broadcast

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

將張量廣播到指定的 GPU 裝置。

引數

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

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

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

注意

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

返回

• 如果指定了 devices，

  返回一個元組，包含 tensor 的副本，這些副本位於 devices 上。

• 如果指定了 out，

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