Merge pull request #3480 from pycaret/fix_tree_plot

pycaret · May 1, 2023 · 9f1e37e · 9f1e37e
2 parents 358f2dd + 5ad349b
commit 9f1e37e
Showing 1 changed file with 7 additions and 8 deletions.
diff --git a/pycaret/internal/pycaret_experiment/tabular_experiment.py b/pycaret/internal/pycaret_experiment/tabular_experiment.py
@@ -1505,15 +1505,14 @@ def tree():
                     self.logger.info("Plotting decision trees")
                     trees = []
                     feature_names = list(self.X_train_transformed.columns)
+                    class_names = None
                     if self._ml_usecase == MLUsecase.CLASSIFICATION:
-                        class_names = {
-                            i: class_name
-                            for i, class_name in enumerate(
-                                get_label_encoder(self.pipeline).classes_
-                            )
-                        }
-                    else:
-                        class_names = None
+                        label_encoder = get_label_encoder(self.pipeline)
+                        if label_encoder:
+                            class_names = {
+                                i: class_name
+                                for i, class_name in enumerate(label_encoder.classes_)
+                            }
                     fitted_estimator = tree_estimator
                     if is_stacked_model:
                         stacked_feature_names = []