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)