torch.get_device_module

torch.get_device_module(device=None) [исходный код]

Возвращает модуль, связанный с указанным устройством (например, torch.device(‘cuda’), “mtia:0”, “xpu”, …). Если устройство не указано, возвращает модуль для текущего ускорителя или для CPU, если ускоритель отсутствует.

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

module