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 的一个副本。