How to use MLlfow to load the logged/saved model in Azure ML?

I want to deploy the trained ML model via AZURE ML online endppoints.

I have already registered my model on the workspace.

Now I am getting following error when I am trying to load the model using cutome score.py for mlflow.pyfunc.load_model()

This is my code -

model_path = os.path.join(os.getenv("AZUREML_MODEL_DIR"), "use-case1-model")
model = mlflow.pyfunc.load_model(model_path)

score.py

import logging
import os
import json
import mlflow
from io import StringIO
from mlflow.pyfunc.scoring_server import infer_and_parse_json_input, predictions_to_json
import sys
from time import strftime, localtime
from collections import Counter
from pytorch_transformers import BertTokenizer
import random
import numpy as np 
import torch 
from tqdm import tqdm

def init():
    global model
    # "model" is the path of the mlflow artifacts when the model was registered. For automl
    # models, this is generally "mlflow-model".
    model_path = os.path.join(os.getenv("AZUREML_MODEL_DIR"), "use-case1-model")
    model = mlflow.pyfunc.load_model(model_path)
    logging.info("Init complete")


def run(raw_data):
    data = json.loads(raw_data)
    title = json.dumps(data["title"])
    att = json.dumps(data["attributes"])

    output = model.predict([tensor_t,tensor_a])

    predict_list = output.tolist()[0]
    
    result = StringIO()
    predictions_to_json(predict_list,result)
    return result.getvalue()

Error that I am getting -

File "/azureml-envs/azureml_9a3b1e0a66d72d612aebc12b4a285f72/lib/python3.9/site-packages/azureml_inference_server_http/server/user_script.py", line 117, in invoke_init
    self._user_init()
  File "/var/azureml-app/dependencies/score.py", line 21, in init
    model = mlflow.pyfunc.load_model(model_path)
  File "/azureml-envs/azureml_9a3b1e0a66d72d612aebc12b4a285f72/lib/python3.9/site-packages/mlflow/pyfunc/__init__.py", line 735, in load_model
    model_impl = importlib.import_module(conf[MAIN])._load_pyfunc(data_path)
  File "/azureml-envs/azureml_9a3b1e0a66d72d612aebc12b4a285f72/lib/python3.9/site-packages/mlflow/pytorch/__init__.py", line 735, in _load_pyfunc
    return _PyTorchWrapper(_load_model(path, **kwargs))
  File "/azureml-envs/azureml_9a3b1e0a66d72d612aebc12b4a285f72/lib/python3.9/site-packages/mlflow/pytorch/__init__.py", line 643, in _load_model
    return torch.load(model_path, **kwargs)
  File "/azureml-envs/azureml_9a3b1e0a66d72d612aebc12b4a285f72/lib/python3.9/site-packages/torch/serialization.py", line 809, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "/azureml-envs/azureml_9a3b1e0a66d72d612aebc12b4a285f72/lib/python3.9/site-packages/torch/serialization.py", line 1172, in _load
    result = unpickler.load()
  File "/azureml-envs/azureml_9a3b1e0a66d72d612aebc12b4a285f72/lib/python3.9/site-packages/torch/serialization.py", line 1142, in persistent_load
    typed_storage = load_tensor(dtype, nbytes, key, _maybe_decode_ascii(location))
  File "/azureml-envs/azureml_9a3b1e0a66d72d612aebc12b4a285f72/lib/python3.9/site-packages/torch/serialization.py", line 1116, in load_tensor
    wrap_storage=restore_location(storage, location),
  File "/azureml-envs/azureml_9a3b1e0a66d72d612aebc12b4a285f72/lib/python3.9/site-packages/torch/serialization.py", line 217, in default_restore_location
    result = fn(storage, location)
  File "/azureml-envs/azureml_9a3b1e0a66d72d612aebc12b4a285f72/lib/python3.9/site-packages/torch/serialization.py", line 182, in _cuda_deserialize
    device = validate_cuda_device(location)
  File "/azureml-envs/azureml_9a3b1e0a66d72d612aebc12b4a285f72/lib/python3.9/site-packages/torch/serialization.py", line 166, in validate_cuda_device
    raise RuntimeError('Attempting to deserialize object on a CUDA '
RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False. If you are running on a CPU-only machine, please use torch.load with map_location=torch.device('cpu') to map your storages to the CPU.

How and where can I update map_location=torch.device('cpu') ? mlflow.pyfunc.load_model() doesnt have a parameter to access map_location and as the packages is installed in docker i cannot make changes to serilaization.py

As per the error logs, you are attempting to deserialize an object on a CUDA device, but torch.cuda.is_available() is returning False, which is due to running on a CPU only machine. To resolve this issue, you need to update the torch.load function to specify map_location=torch.device('cpu') to map the storages to the CPU.

Since the mlflow.pyfunc.load_model() function does not have a map_location argument, you can use a **kwargs argument that can pass any additional keyword arguments to the torch.load() function.

To solve the issue, add *{'map_location': torch.device('cpu')}in your score.py file.

model = mlflow.pyfunc.load_model(model_path, *{'map_location': torch.device('cpu')}) 

or Use below code:(Updated solution)

model = mlflow.pytorch.load_model(model_path, map_location=torch.device('cpu'))

Example:

import  mlflow
import  torch
path="./deploy/credit_defaults_model/"
model = mlflow.pyfunc.load_model(path, *{'map_location': torch.device('cpu')})