torch.accelerator.current_accelerator

torch.accelerator.current_accelerator(check_available=False)

返回編譯時可用的加速器的裝置。如果沒有在編譯時提供加速器，則返回 None。有關詳細資訊，請參閱 accelerator。

引數

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

返回

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

返回型別

torch.device

注意

返回的 torch.device 的索引將為 None，請使用 torch.accelerator.current_device_index() 來了解當前正在使用的索引。此 API 不會汙染 fork。有關更多詳細資訊，請參閱 多程序中的 Poison 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)