mlflow/mlflow · error · MlflowException

`tensorflow` must be installed if you want to load an export

Error message

`tensorflow` must be installed if you want to load an exported Keras 3 model, please install `tensorflow` by `pip install tensorflow`.

What it means

When a Keras 3 model was saved with save_exported_model=True, it was exported as a TensorFlow SavedModel directory. Loading such a model via mlflow.keras.load_model or the pyfunc loader requires the `tensorflow` package to deserialize tf.saved_model.load; if tensorflow is not importable in the current environment, an MlflowException is raised with install instructions.

Source

Thrown at mlflow/keras/load.py:60

            raise MlflowException(
                f"`data` must be one of: {[x.__name__ for x in supported_input_types]}, but "
                f"received type: {type(data)}.",
                INVALID_PARAMETER_VALUE,
            )
        # Return numpy array for serving purposes.
        return keras.ops.convert_to_numpy(model_call(data))


def _load_keras_model(path, model_conf, custom_objects=None, **load_model_kwargs):
    save_exported_model = model_conf.flavors["keras"].get("save_exported_model")
    model_path = os.path.join(path, model_conf.flavors["keras"].get("data", _MODEL_SAVE_PATH))
    if os.path.isdir(model_path):
        model_path = os.path.join(model_path, _MODEL_SAVE_PATH)
    if save_exported_model:
        try:
            import tensorflow as tf
        except ImportError:
            raise MlflowException(
                "`tensorflow` must be installed if you want to load an exported Keras 3 model, "
                "please install `tensorflow` by `pip install tensorflow`."
            )
        return tf.saved_model.load(model_path)
    else:
        model_path += ".keras"
        return keras.saving.load_model(
            model_path,
            custom_objects=custom_objects,
            **load_model_kwargs,
        )


def load_model(model_uri, dst_path=None, custom_objects=None, load_model_kwargs=None):
    """
    Load Keras model from MLflow.

    This method loads a saved Keras model from MLflow, and returns a Keras model instance.

View on GitHub (pinned to 6a27f2decc)

Solutions

Install TensorFlow in the current environment: pip install tensorflow
Or pip install 'mlflow[keras]' / the extra that pulls tensorflow
If you do not need the exported SavedModel, re-log the model with save_exported_model=False and load the .keras artifact instead

Example fix

// before
pip install mlflow
model = mlflow.keras.load_model("runs:/abc/model")
// after
pip install mlflow tensorflow
model = mlflow.keras.load_model("runs:/abc/model")

Defensive patterns

Strategy: validation

Validate before calling

import importlib.util
def ensure_tensorflow():
    if importlib.util.find_spec("tensorflow") is None:
        raise RuntimeError("pip install tensorflow before loading this exported Keras 3 model")

Type guard

def tensorflow_available() -> bool:
    import importlib.util
    return importlib.util.find_spec("tensorflow") is not None

Try / catch

from mlflow.exceptions import MlflowException
try:
    model = mlflow.keras.load_model(model_uri)
except MlflowException as e:
    import subprocess; subprocess.run(["pip", "install", "tensorflow"], check=True)
    model = mlflow.keras.load_model(model_uri)

Prevention

Pin tensorflow in serving/deployment images when models use save_exported_model=True
Check MLMmodel flavor metadata for save_exported_model before choosing the runtime
Use 'pip install mlflow[keras]' to get matching deps

When it happens

Trigger: Calling mlflow.keras.load_model(model_uri) (or serving via _load_pyfunc) on a logged Keras 3 model whose MLMODEL flavor metadata contains save_exported_model=True, in an environment where `import tensorflow` fails (tensorflow not installed or broken install).

Common situations: Deploying to a slim serving image that only has keras/torch deps and not full tensorflow; skinny client installs; loading a model logged on a machine with TF into a TF-free environment.

Related errors

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