torch.cuda.comm.broadcast_coalesced

torch.cuda.comm.broadcast_coalesced(tensors, devices, buffer_size=10485760) [исходный код]

Рассылает последовательность тензоров на указанные GPU.

Сначала небольшие тензоры объединяются в буфер, чтобы сократить число синхронизаций.

Параметры:

• tensors (sequence) – тензоры для рассылки. Все они должны находиться на одном устройстве: CPU или GPU.
• devices (Iterable[torch.device, str or int]) – итерируемый объект с GPU-устройствами, между которыми выполняется рассылка.
• buffer_size (int) – максимальный размер буфера, используемого для объединения

Возвращает:

Кортеж, содержащий копии tensor, размещённые на devices.