INNER CODE UNIT · Python

_to_device

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

def _to_device(model: Any, args: argparse.Namespace) -> Any:
    """Load the predictor on the requested device

    Args:
        model: the predictor
        args: the parsed command-line arguments

    Returns:
        the predictor, loaded on the requested device
    """
    device = _resolve_device(args.device)
    try:
        model = model.to(device)
        if device.type == "cuda" and torch.cuda.get_device_capability(device) >= (8, 0):
            model = model.bfloat16()
            logger.info(f"Model loaded on {device} with bfloat16 precision")
    except (RuntimeError, AssertionError, ValueError) as e:
        logger.error(f"Could not load the model on device '{device}': {e}")

View source record →