device_index

class torch.accelerator.device_index(device, /)

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

参数

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