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)