torch.accelerator.current_accelerator

torch.accelerator.current_accelerator(check_available=False)

返回編譯時可用的加速器裝置。如果編譯時沒有可用的加速器，則返回 None。詳見 accelerator。

引數

check_available (bool, 可選) – 如果為 True，除了編譯時檢查外，還將進行執行時檢查以檢視裝置 torch.accelerator.is_available() 是否可用。預設值：False

返回

將當前加速器作為 torch.device 返回。

返回型別

torch.device

注意

返回的 torch.device 的索引將為 None，請使用 torch.accelerator.current_device_index() 來獲取當前使用的索引。

示例

>>> # 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)