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 张量，每个张量都包含 tensor 的副本。