[Frameworks] Fix autolog column name detection (#1292)

mlrun · Sep 5, 2021 · da6d288 · da6d288
1 parent 5479807
commit da6d288
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/mlrun/frameworks/mlbase/__init__.py b/mlrun/frameworks/mlbase/__init__.py
@@ -0,0 +1 @@
+from .._common.pkl_model_server import PickleModelServer  # noqa
diff --git a/mlrun/frameworks/mlbase/mlrun_interface.py b/mlrun/frameworks/mlbase/mlrun_interface.py
@@ -74,6 +74,11 @@ def _post_fit(*args, **kwargs):
                     )
 
             # Log fitted model and metrics
+            label_column = (
+                y_train.name
+                if isinstance(y_train, pd.Series)
+                else y_train.columns.to_list()
+            )
             context.log_model(
                 model_name or "model",
                 db_key=model_name,
@@ -85,6 +90,6 @@ def _post_fit(*args, **kwargs):
                 metrics=context.results,
                 format="pkl",
                 training_set=train_set,
-                label_column=y_train.columns.to_list(),
+                label_column=label_column,
                 extra_data=eval_metrics,
             )