INNER CODE UNIT · Python

_resolve_device

mindee/doctr · doctr/cli/main.py:129

def _resolve_device(device: str | None) -> torch.device:
    """Resolve the device the models should be loaded on

    Args:
        device: the `--device` value

    Returns:
        the resolved torch device
    """
    if device is None or device.strip().lower() == "auto":
        return torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    try:
        return torch.device(device)
    except (RuntimeError, ValueError) as e:
        logger.error(f"Invalid device '{device}': {e}")
        sys.exit(1)

View source record →