def _load_DINO_model(arch: Literal['vit_small', 'vit_base'], device: str, patch_size: int = None,
                    image_size: Tuple = (224, 224)) -> Tuple[Any, Any]:
    """ loads a vit dino model
       # arch -> ['vit_tiny', 'vit_small', 'vit_base']
    Args:
        arch (str): the model architecture
        device (str): cpu or cuda
        patch_size (_type_, optional): the number of patches for the model. note this needs to exactly match
                                        the number the model was trained with.

    Raises:
        KeyError: _description_
        KeyError: _description_
        RuntimeError: _description_

    Returns:
        Tuple[Any, Any]: the model and transform for the inputs
    """
  
    allowed_archs = ('vit_small', 'vit_base')
    allowed_patches = (8, 16)
    if arch not in allowed_archs:
        raise KeyError(f"{arch} not found in {allowed_archs}")

    if patch_size not in allowed_patches:
        raise KeyError(f"{patch_size} not found in {allowed_patches}")

    if arch == "vit_small":
        # TODO add onnx support
        if patch_size == 8:
            model = torch.hub.load('facebookresearch/dino:main', 'dino_vits8')

        if patch_size == 16:
            model = torch.hub.load('facebookresearch/dino:main', 'dino_vits16')

    elif arch == "vit_base":
        if patch_size == 8:
            model = torch.hub.load('facebookresearch/dino:main', 'dino_vitb8')

        if patch_size == 16:
            model = torch.hub.load('facebookresearch/dino:main', 'dino_vitb16')

    else:
        raise ModelLoadError(f"unknown model of arch:{arch} and patch_size:{patch_size}")

    model.eval()
    model.to(device)

    return model, _get_DINO_transform(image_size=image_size)
 loads a vit dino model
       # arch -> ['vit_tiny', 'vit_small', 'vit_base']
    Args:
        arch (str): the model architecture
        device (str): cpu or cuda
        patch_size (_type_, optional): the number of patches for the model. note this needs to exactly match
                                        the number the model was trained with.

    Raises:
        KeyError: _description_
        KeyError: _description_
        RuntimeError: _description_

    Returns:
        Tuple[Any, Any]: the model and transform for the inputs
    
src/marqo/s2_inference/processing/DINO_utils.py