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