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。有关详细信息，请参阅 多进程中的毒害 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)