Initial update to load model

Updated load model step, initial changes.
Nixtla · Oct 12, 2024 · ae6db6e · ae6db6e
1 parent 62d10a8
commit ae6db6e
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/neuralforecast/flavor.py b/neuralforecast/flavor.py
@@ -278,7 +278,7 @@ def log_model(
 
 def load_model(model_uri, dst_path=None):
     """
-    Load an ``MLForecast`` model from a local file or a run.
+    Load an ``NeuralForecast`` model from a local file or a run.
 
     Parameters
     ----------
@@ -301,13 +301,13 @@ def load_model(model_uri, dst_path=None):
 
     Returns
     -------
-    An ``MLForecast`` model instance.
+    An ``NeuralForecast`` model instance.
     """
     local_model_path = _download_artifact_from_uri(artifact_uri=model_uri, output_path=dst_path)
     flavor_conf = _get_flavor_configuration(model_path=local_model_path, flavor_name=FLAVOR_NAME)
     _add_code_from_conf_to_system_path(local_model_path, flavor_conf)
     model_file_path = os.path.join(local_model_path, flavor_conf["pickled_model"])
-    return MLForecast.load(model_file_path)
+    return NeuralForecast.load(model_file_path)
 
 
 def _load_pyfunc(path):