torch.accelerator.current_accelerator

torch.accelerator.current_accelerator(check_available=False)

返回编译时可用的加速器设备。如果编译时没有加速器可用，则返回 None。有关详细信息，请参阅 加速器。

参数

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)