torch.accelerator.get_device_capability

torch.accelerator.get_device_capability(device=None, /) [исходный код]

Возвращает возможности текущего выбранного устройства.

Параметры:

device (torch.device, str, int, необязательно) – Тип устройства ускорителя, для которого нужно запросить возможности. Если не задано, по умолчанию используется torch.accelerator.current_device_index().

Возвращает:

Словарь со сведениями о возможностях устройства. Словарь содержит:

• supported_dtypes (set(torch.dtype)): Набор типов данных PyTorch, для которых на ускорителе можно выделять память под тензоры и поддерживается преобразование типов между поддерживаемыми типами данных. Поддержка любых операторов за пределами этого набора не гарантируется

Тип возвращаемого значения:

dict[str, Any]

Примеры

>>> # Query capabilities for current device
>>> capabilities = torch.accelerator.get_device_capability("cuda:0")
>>> print("Supported dtypes:", capabilities["supported_dtypes"])