Merge pull request #90 from qrsforever/master

fix bug: check model attrs
parrt · Jun 4, 2020 · b38b3cc · b38b3cc
2 parents be7fcce + 2154371
commit b38b3cc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/dtreeviz/shadow.py b/dtreeviz/shadow.py
@@ -40,7 +40,7 @@ def __init__(self, tree_model,
         self.class_names = class_names
         self.class_weight = tree_model.class_weight
 
-        if getattr(tree_model, 'tree_') is None: # make sure model is fit
+        if not hasattr(tree_model, 'tree_'): # make sure model is fit
             tree_model.fit(X_train, y_train)
 
         if tree_model.tree_.n_classes > 1: