torch.cuda.comm.gather

torch.cuda.comm.gather(tensors, dim=0, destination=None, *, out=None) [source]

Собирает тензоры с нескольких устройств GPU.

Параметры

• tensors (Iterable[Tensor]) – итерируемый набор тензоров для сбора. Размеры тензоров во всех измерениях, кроме dim, должны совпадать.
• dim (int, необязательно) – измерение, по которому будут конкатенированы тензоры. По умолчанию: 0.
• destination (torch.device, str, или int, необязательно) – устройство вывода. Может быть CPU или CUDA. По умолчанию: текущее устройство CUDA.
• out (Tensor, необязательно, только ключевое слово) – тензор для хранения результата сбора. Его размеры должны совпадать с размерами tensors, за исключением dim, где размер должен быть равен sum(tensor.size(dim) for tensor in tensors). Может находиться на CPU или CUDA.

Примечание

destination не должен быть указан, когда указан out.

Возвращает

• If destination is specified,

  тензор, расположенный на устройстве destination, являющийся результатом конкатенации tensors по dim.

• If out is specified,

  тензор out, теперь содержащий результаты конкатенации tensors по dim.