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