espnet2.torch_utils.resolve_device.resolve_device

Less than 1 minute

espnet2.torch_utils.resolve_device.resolve_device

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 CUDA when torch sees a CUDA device, else the CPU, so a caller need not know whether this machine has a GPU. 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", "cuda" or "cuda:<n>".
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 a GPU
'cpu'
>>> resolve_device("cuda:0")    # on one with a GPU
'cuda:0'
>>> resolve_device("gpu")
Traceback (most recent call last):
ValueError: unknown device 'gpu': use auto, cpu, mps, cuda or cuda:&lt;n&gt;