mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 12:54:11 +08:00
1239260a0e
As the title said, `acc` already exists so we just use it instead of calling `current_accelerator()` again. Pull Request resolved: https://github.com/pytorch/pytorch/pull/150829 Approved by: https://github.com/guangyey, https://github.com/Skylion007
29 lines
968 B
Python
29 lines
968 B
Python
from typing import Optional
|
|
|
|
import torch
|
|
from torch.types import Device as _device_t
|
|
|
|
|
|
def _get_device_index(device: _device_t, optional: bool = False) -> int:
|
|
if isinstance(device, int):
|
|
return device
|
|
if isinstance(device, str):
|
|
device = torch.device(device)
|
|
device_index: Optional[int] = None
|
|
if isinstance(device, torch.device):
|
|
acc = torch.accelerator.current_accelerator()
|
|
if acc is None:
|
|
raise RuntimeError("Accelerator expected")
|
|
if acc.type != device.type:
|
|
raise ValueError(
|
|
f"{device.type} doesn't match the current accelerator {acc}."
|
|
)
|
|
device_index = device.index
|
|
if device_index is None:
|
|
if not optional:
|
|
raise ValueError(
|
|
f"Expected a torch.device with a specified index or an integer, but got:{device}"
|
|
)
|
|
return torch.accelerator.current_device_index()
|
|
return device_index
|