Remove tree_index from spark shadow tree initialization

parrt · Feb 7, 2021 · 78d89d9 · 78d89d9
1 parent c20c257
commit 78d89d9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/dtreeviz/models/shadow_decision_tree.py b/dtreeviz/models/shadow_decision_tree.py
@@ -464,7 +464,7 @@ def get_shadow_tree(tree_model, x_data, y_data, feature_names, target_name, clas
         elif (str(type(tree_model)).endswith("pyspark.ml.classification.DecisionTreeClassificationModel'>") or
               str(type(tree_model)).endswith("pyspark.ml.classification.DecisionTreeClassificationModel'>")):
             from dtreeviz.models import spark_decision_tree
-            return spark_decision_tree.ShadowSparkTree(tree_model, tree_index, x_data, y_data,
+            return spark_decision_tree.ShadowSparkTree(tree_model, x_data, y_data,
                                                        feature_names, target_name, class_names)
         elif "lightgbm.basic.Booster" in str(type(tree_model)):
             from dtreeviz.models import lightgbm_decision_tree