torch.cuda.comm.scatter

torch.cuda.comm.scatter(tensor, devices=None, chunk_sizes=None, dim=0, streams=None, *, out=None) [source]

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

Параметры

• tensor (Tensor) – тензор для рассылки. Может находиться на CPU или GPU.
• devices (Iterable[torch.device, str или int], необязательный) – итерируемый список устройств GPU, по которым следует произвести рассылку.
• chunk_sizes (Iterable[int], необязательный) – размеры кусков, которые будут размещены на каждом устройстве. Длина должна соответствовать devices, а сумма должна составлять tensor.size(dim). Если не указано, то tensor будет разделен на равные куски.
• dim (int, необязательный) – измерение, по которому следует разбить tensor. Значение по умолчанию: 0.
• streams (Iterable[torch.cuda.Stream], необязательный) – итерируемый список потоков, в которых следует выполнить рассылку. Если не указано, будет использован потоком по умолчанию.
• out (Sequence[Tensor], необязательный, только ключевое слово) – тензоры GPU для хранения результатов. Размер этих тензоров должен соответствовать размеру tensor, за исключением dim, где общая сумма должна составлять tensor.size(dim).

Примечание

Должен быть указан ровно один из devices и out. Когда out указан, chunk_sizes не должен быть указан и будет вычислен из размеров out.

Возвращает

• If devices is specified,

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

• If out is specified,

  кортеж, содержащий out тензоров, каждый из которых содержит кусок tensor.