torch.cuda.comm.broadcast_coalesced

torch.cuda.comm.broadcast_coalesced(tensors, devices, buffer_size=10485760)

将一系列张量广播到指定的GPU。

小张量首先被合并到一个缓冲区中，以减少同步的数量。

Parameters

• 张量 (序列) – 要广播的张量。必须位于同一设备上，无论是CPU还是GPU。

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

• buffer_size (int) – 用于合并的最大缓冲区大小

Returns

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