torch.accelerator.current_accelerator

torch.accelerator.current_accelerator(check_available=False) [исходный код]

Возвращает устройство ускорителя, доступного на этапе компиляции. Если на этапе компиляции ускоритель недоступен, возвращает None. Подробнее см. в разделе ускорители.

Параметры:

check_available (bool, необязательно) – если True, дополнительно выполняется проверка во время выполнения, чтобы определить, доступно ли устройство (torch.accelerator.is_available()), помимо проверки на этапе компиляции. По умолчанию: False

Возвращает:

текущий ускоритель в виде torch.device.

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

torch.device

Примечание

Индекс возвращаемого torch.device будет равен None; чтобы узнать текущий используемый индекс, воспользуйтесь torch.accelerator.current_device_index(). Этот API НЕ делает fork небезопасным. Подробнее см. в разделе Небезопасное использование fork при многопроцессной обработке.

Пример:

>>> # If an accelerator is available, sent the model to it
>>> model = torch.nn.Linear(2, 2)
>>> if (current_device := current_accelerator(check_available=True)) is not None:
>>>     model.to(current_device)