device_index

class torch.accelerator.device_index(device, /)

用于设置当前 加速器 的当前设备索引的上下文管理器。它会在上下文的持续时间内将当前设备索引临时更改为指定值，并在退出上下文时自动恢复之前的设备索引。

参数

device (Optional[int]) – 要临时设置的设备索引。如果为 None，则不进行设备索引切换。

示例

>>> # Set device 0 as the current device temporarily
>>> with torch.accelerator.device_index(0):
...     # Code here runs with device 0 as the current device
...     pass
>>> # Original device is now restored
>>> # No-op when None is passed
>>> with torch.accelerator.device_index(None):
...     # No device switching occurs
...     pass