device_index#

classtorch.accelerator.device_index(device,/)[source]#

Context manager to set the current device index for the currentaccelerator.Temporarily changes the current device index to the specified value for the durationof the context, and automatically restores the previous device index when exitingthe context.

Parameters

device (Optional[int]) – a given device index to temporarily set. If None,no device index switching occurs.

Examples

>>># Set device 0 as the current device temporarily>>>withtorch.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>>>withtorch.accelerator.device_index(None):...# No device switching occurs...pass