torch.cuda.comm.gather

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

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

Параметры:

• tensors (Iterable[Tensor]) – итерируемый объект с тензорами для сбора. Размеры тензоров по всем измерениям, кроме dim, должны совпадать.
• dim (int, optional) – измерение, вдоль которого будут объединены тензоры. Значение по умолчанию: 0.
• destination (torch.device, str, or int, optional) – устройство для выходных данных. Может быть CPU или CUDA. Значение по умолчанию: текущее устройство CUDA.
• out (Tensor, optional, keyword-only) – тензор для сохранения результата сбора. Его размеры должны совпадать с размерами 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.