mlflow/mlflow · error · MlflowException

Couldn't find a loader class for {class_name}

Error message

Couldn't find a loader class for {class_name}

What it means

MLflow's transformers flavor supports loading custom models defined via a config's auto_map. When building the loader, _load_class_from_transformers_config scans config.auto_map for an entry whose module path ends with the requested class name; if none matches, it cannot resolve a loader class and raises this MlflowException.

Source

Thrown at mlflow/transformers/model_io.py:290

    # if the class is available in transformers natively,
    # then we don't need to execute any custom code.
    if hasattr(transformers, class_name):
        cls = getattr(transformers, class_name)
        return cls, False
    else:
        # else, we need to fetch the correct AutoClass.
        # this is defined in the `auto_map` field. there
        # should only be one AutoClass that maps to the
        # model's class name.
        auto_classes = [
            auto_class
            for auto_class, module in config.auto_map.items()
            if module.split(".")[-1] == class_name
        ]

        if len(auto_classes) == 0:
            raise MlflowException(f"Couldn't find a loader class for {class_name}")

        auto_class = auto_classes[0]
        cls = getattr(transformers, auto_class)

        # we will need to trust remote code when loading the model
        return cls, True


def _load_model(model_name_or_path, flavor_conf, accelerate_conf, device, revision=None):
    """
    Try to load a model with various loading strategies.
      1. Try to load the model with accelerate
      2. Try to load the model with the specified device
      3. Load the model without the device
    """
    import transformers

    if hasattr(transformers, flavor_conf[FlavorKey.MODEL_TYPE]):

View on GitHub (pinned to 6a27f2decc)

Solutions

1. Open the model's config.json and ensure auto_map includes an entry whose module path ends with the class name being loaded (e.g. 'AutoModelForSeq2SeqLM': 'modeling.CustomModel').
2. Add the missing auto_map entry and re-save/re-upload the model.
3. Load the model directly with transformers (AutoModel.from_pretrained(..., trust_remote_code=True)) to verify the config is valid before going through MLflow.
4. Ensure you are loading the intended revision/commit of the repo — an older revision may lack the auto_map entry.

Example fix

// before (config.json)
{"auto_map": {"AutoModel": "modeling_chat.CustomChatModel"}}
// after
{"auto_map": {"AutoModel": "modeling_chat.CustomChatModel", "AutoModelForCausalLM": "modeling_chat.CustomChatModel"}}

import json
with open(f'{model_dir}/config.json') as f:
    cfg = json.load(f)
assert any(k.endswith(class_name) or v.split('.')[-1] == class_name for k, v in cfg.get('auto_map', {}).items()), f"auto_map lacks entry for {class_name}"

def has_auto_map_entry(config, class_name: str) -> bool:
    am = getattr(config, 'auto_map', None) or {}
    return any(module.split('.')[-1] == class_name for module in am.values())

try:
    model = mlflow.transformers.load_model(uri)
except MlflowException as e:
    if "Couldn't find a loader class" in str(e):
        model = AutoModel.from_pretrained(path, trust_remote_code=True)
    else:
        raise

When it happens

Trigger: Loading (mlflow.transformers.load_model or pyfunc load) a custom transformers model whose config.json auto_map does not contain a key mapping to a module ending with the expected class name (e.g. AutoModelForSeq2SeqLM missing), or the auto_map entries point at different class names than the saved model.

Common situations: Custom 'trust_remote_code' repos where the author renamed classes or only registered some of the auto_map entries; hand-edited or incomplete config.json; loading a model saved with a different transformers architecture than expected.

Related errors

• Failed to load base model '{effective_base_model}'. If the m
• Failed to load base model '{base_model}'. If the model has m
• Loading {config_type} chain not supported
• INVALID_STATE
• The model was saved with a HuggingFace Hub repository name '

AI-assisted analysis of mlflow/mlflow@6a27f2decc (2026-08-29). Data as JSON: /api/errors/fd03f7b67b55b120. Report an issue: GitHub.