Fix tuning of stacked models

pycaret · Jan 22, 2021 · 2b3ae9c · 2b3ae9c
1 parent 236bc3e
commit 2b3ae9c
Showing 1 changed file with 17 additions and 7 deletions.
diff --git a/pycaret/internal/tabular.py b/pycaret/internal/tabular.py
@@ -4041,7 +4041,14 @@ def tune_model_supervised(
 
     logger.info("Checking base model")
 
-    estimator_id = _get_model_id(estimator)
+    is_stacked_model = False
+
+    if hasattr(estimator, "final_estimator"):
+        logger.info("Model is stacked, using the definition of the meta-model")
+        is_stacked_model = True
+        estimator_id = _get_model_id(estimator.final_estimator)
+    else:
+        estimator_id = _get_model_id(estimator)
     if estimator_id is None:
         if custom_grid is None:
             raise ValueError(
@@ -4059,15 +4066,18 @@ def tune_model_supervised(
         model = clone(estimator)
     else:
         logger.info("Model has a special tunable class, using that")
-        model = clone(estimator_definition.tunable(**estimator.get_params()))
-    is_stacked_model = False
+        if is_stacked_model:
+            model = clone(estimator)
+            model.set_params(
+                final_estimator=estimator_definition.tunable(**estimator.get_params())
+            )
+        else:
+            model = clone(estimator_definition.tunable(**estimator.get_params()))
 
     base_estimator = model
 
-    if hasattr(base_estimator, "final_estimator"):
-        logger.info("Model is stacked, using the definition of the meta-model")
-        is_stacked_model = True
-        base_estimator = base_estimator.final_estimator
+    if is_stacked_model:
+        base_estimator = model.final_estimator
 
     display.update_monitor(2, estimator_name)
     display.display_monitor()