torch.cuda.comm.broadcast

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

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

参数:

• 张量（Tensor）- 要广播的张量。可以在 CPU 或 GPU 上。

• devices（可迭代对象[torch.device, str 或 int]，可选）- 在其中广播的 GPU 设备集合。

• out（序列[Tensor]，可选，关键字参数）- 存储输出结果的 GPU 张量。

注意

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

返回值:

• 如果指定了 devices ，

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

• 如果指定了 out ，

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