torch.cuda.comm.reduce_add_coalesced

torch.cuda.comm.reduce_add_coalesced(inputs, destination=None, buffer_size=10485760) [исходный код]

Суммирует тензоры с нескольких графических процессоров.

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

Параметры:

• inputs (Iterable[Iterable[Tensor]]) – итерируемый объект, содержащий итерируемые объекты с тензорами с одного устройства.
• destination (int, optional) – устройство, на котором будет размещен результат (по умолчанию: текущее устройство).
• buffer_size (int) – максимальный размер буфера, используемого для объединения

Возвращает:

Кортеж тензоров, содержащих поэлементную сумму каждой группы входных данных и размещенных на устройстве destination.