espnet2.torch_utils.resolve_device.resolve_device

source

espnet2.torch_utils.resolve_device.resolve_device(requested: str = 'auto') → str

Return the device a model is built on, for "auto" or a named one.

"auto" is the accelerator torch sees – CUDA, NPU, XPU, whichever torch.accelerator reports – else the CPU, so a caller need not know what this machine has. MPS is not chosen on its own: it is faster than the CPU for some models’ encoders and slower for a CTC prefix search, and until that is measured per model it is used only when asked for by name.

A named device is checked against what torch has, so a typo or a device this machine lacks fails here, before a model is downloaded and built, rather than in the first forward pass.

• Parameters:requested – "auto", "cpu", "mps", or an accelerator name with or without an index, e.g. "cuda", "cuda:1", "npu", "npu:3".
• Returns: never "auto".
• Return type: A device torch accepts
• Raises:ValueError – If requested names no device, or one this machine does not have.

Examples

>>> resolve_device()            # a machine without an accelerator
'cpu'
>>> resolve_device("cuda:0")    # on one with a GPU
'cuda:0'
>>> resolve_device("npu:0")     # on one with an Ascend NPU
'npu:0'
>>> resolve_device("gpu")
Traceback (most recent call last):
ValueError: unknown device 'gpu': use auto, cpu, mps, or an ...