torch.cuda.comm.broadcast_coalesced

torch.cuda.comm.broadcast_coalesced(tensors, devices, buffer_size=10485760)[源代码]

将张量序列传输到指定的GPU。

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

参数

• tensors (序列) – 需要广播的张量。所有张量必须位于同一设备上，即要么都在CPU上，要么都在GPU上。

• devices (Iterable[torch.device, str 或 int]) – 一个包含要进行广播的 GPU 设备的可迭代对象。

• buffer_size (int) – 缓冲区的最大尺寸，用于合并操作

返回值

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