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.