Hi, I'm using Databricks Feature Store to register a custom model using a model wrapper as follows:

# Log custom model to MLflow
fs.log_model(
  artifact_path="model",
  model = production_model,
  flavor = mlflow.pyfunc, 
  training_set = training_set,
  registered_model_name = model_name,
  conda_env=model_env
)

The model is correctly serving and I can score batch using the model's uri:

fs.score_batch('models:/model_name/latest',
                dataset,
                result_type = ArrayType(StringType()))

What I can not do is to load the model using: 

import mlflow
import databricks.feature_store
mlflow.pyfunc.load_model('models:/model_name/latest')

 I keep on getting the following error: 

ModuleNotFoundError: No module named 'databricks.feature_store.mlflow_model'

What is the correct way to load or log the model for this case?