torch.cuda.comm.broadcast

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

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

Parameters

• 张量 (张量) – 要广播的张量。可以在CPU或GPU上。

• 设备（可迭代[torch.device, str 或 int], 可选）——一个包含GPU设备的可迭代对象，在这些设备之间进行广播。

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

注意

必须指定 devices 和 out 中的一个。

Returns

• 如果指定了devices，

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

• 如果指定了out，

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