ModuleNotFoundError: No module named 'model_train' when using mlflow.sklearn.load_model

Machine Learning

Dive into the world of machine learning on the Databricks platform. Explore discussions on algorithms, model training, deployment, and more. Connect with ML enthusiasts and experts.

Hello,

I have multiple versions of a model registered in model registry. When I am trying to load any other version except model version 1 by

mlflow.sklearn.load_model(f"models:/{model_name}/{model_version}")

I am getting ModuleNotFoundError: No module named 'model_train'.

All the versions have been registered successfully, below is the log when the latest version was registered.

Registered model <model_name> already exists. Creating a new version of this model...

2024/08/22 13:22:58 INFO mlflow.store.model_registry.abstract_store: Waiting up to 300 seconds for model version to finish creation. Model name: <model_name>, version 7

Created version '7' of model <model_name>.