Error al cargar el modelo de MLFlow en Python 3.7

Estoy guardando el modelo de MLFlow usando Databricks. A continuación se detallan los detalles:

ruta_de_artefacto: modelo
databricks_runtime: 8.4.x-gpu-ml-scala2.12
sabores:
python_function:
data: datos
env: conda.yaml
loader_module: mlflow.pytorch
pickle_module_name: mlflow.pytorch.pickle_module
versión_de_python: 3.8.8
pytorch:
datos_del_modelo: datos
versión_de_pytorch: 1.9.0+cu102

No puedo cargar el modelo localmente usando Python 3.7, mientras que funciona bien con Python 3.9.

¿Alguna idea de cuál podría ser la posible resolución a esto?

Rastro de error:

Archivo “/Users/danishm/opt/miniconda3/envs/py37/lib/python3.7/site-packages/mlflow/pytorch/init.py”, línea 714, en cargar_modelo
return _load_model(path=torch_model_artifacts_path, **kwargs)
Archivo “/Users/danishm/opt/miniconda3/envs/py37/lib/python3.7/site-packages/mlflow/pytorch/__init__.py”, línea 626, en _load_model
return torch.load(model_path, **kwargs)
Archivo “/Users/danishm/opt/miniconda3/envs/py37/lib/python3.7/site-packages/torch/serialization.py”, línea 607, en load
return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
Archivo “/Users/danishm/opt/miniconda3/envs/py37/lib/python3.7/site-packages/torch/serialization.py”, línea 882, en _load
result = unpickler.load()
TypeError: code() toma como máximo 15 argumentos (16 dados)

He intentado especificar el nombre del módulo pickle explícitamente como mlflow.pytorch.load_model(ML_MODEL_PATH,pickle_module=mlflow.pytorch.pickle_module)

Pero sigo obteniendo el mismo error.