# Load\_model from MLflow model registry can't find model class - how best to make this available?

**URL:** <https://discuss.pytorch.org/t/load-model-from-mlflow-model-registry-cant-find-model-class-how-best-to-make-this-available/215230>\
**Category:** deployment\
**Created:** [January 10, 2025, 7:46pm UTC](https://discuss.pytorch.org/t/load-model-from-mlflow-model-registry-cant-find-model-class-how-best-to-make-this-available/215230 "2025-01-10T19:46:48Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![minamesephy](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/minamesephy/32/75028_2.png) [@minamesephy](https://discuss.pytorch.org/u/minamesephy)\
**Post date:** [January 10, 2025, 7:46pm UTC](https://discuss.pytorch.org/t/load-model-from-mlflow-model-registry-cant-find-model-class-how-best-to-make-this-available/215230/1 "2025-01-10T19:46:48Z")

</div>

I am working to support some model developers by prototyping some functionality with MLflow model registry.

We successfully register versions of the model fine (it’s a Pytorch Lightning model), and I recently also added code\_paths to the log\_model call so that theoretically all of the original model code ends up registered / in the artifacts (I’m not sure if this is best practice though? But either way the model code is super complex and I’m not sure how I can extract out just the stuff needed by the serializer on the consuming side atm).

```auto
        mlflow.pytorch.log_model(
            pytorch_model=cli.model,
            artifact_path="model",
            registered_model_name="MyModelName",
            code_paths=["src"],
        )

```

In any case, I can see the model registered and the artifacts tree looks like this now when I view them from the UI:

```auto
model
|--checkpoints
|--code
    |--src
        |--<mymodelname>
             |-- a bunch of other modules and nested stuff under here
        |--a bunch of other stuff here
|--data
     |--model.pth
     |--pickle_module_info.txt (this is a small file with just the text "mlflow.pytorch.pickle_module" in it)
|--MLmodel (small file w/ metadata about the model)
|--conda.yaml
|--python_env.yaml
|--requirements.txt

```

I have a totally separate package outside of the model package and I’m testing consuming the model as an “external” client who wants to use the registered pre-trained model for further fine-tuning or inference. I confirmed that the code is getting retrieved when I do something like this:

```auto
model_uri = f"models:/{model_name}/{model_version}"
local_path = mlflow.artifacts.download_artifacts(model_uri)

```

the retrieved artifacts there include everything I listed above.

However, when I try to load the model:

```auto
    loaded_model = mlflow.pytorch.load_model(model_uri)

```

It fails with the error: `ModuleNotFoundError: No module named 'src.<mymodelname>'`

Am I going about this totally wrong or is there just something small I’m missing here (like getting the picker to recognize the original model code under code/ somehow since it’s all there?)?

Any help would be much appreciated!