torch.cuda.comm.broadcast

torch.cuda.comm.broadcast(tensor, devices=None, *, out=None)[源代码]

将张量广播到指定的GPU设备上。

参数

• tensor (Tensor) – 需要进行广播的张量，可以位于CPU或GPU上。

• devices (Iterable[torch.device, str 或 int], 可选) – 在这些 GPU 设备上进行广播。

• out (Sequence[Tensor], 可选, 关键字参数) – 用于存储输出结果的 GPU 张量。

注意

必须指定 devices 和 out 中的其中一个，但不能同时指定两者。

返回值

• 如果指定 devices，

  一个包含tensor副本的元组，这些副本被放置在devices上。

• 如果指定 out，

  包含out张量的元组，每个out张量都包含了tensor的一个副本。