make _task vs task be version sensitive.

Signed-off-by: Terence Parr <parrt@antlr.org>
parrt · Jan 28, 2023 · 2f10aaa · 2f10aaa
1 parent ffda2a0
commit 2f10aaa
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/dtreeviz/models/tensorflow_decision_tree.py b/dtreeviz/models/tensorflow_decision_tree.py
@@ -2,6 +2,7 @@
 from typing import List, Mapping
 
 import numpy as np
+import tensorflow_decision_forests
 from tensorflow_decision_forests.component.py_tree.node import LeafNode
 from tensorflow_decision_forests.keras import RandomForestModel
 from tensorflow_decision_forests.tensorflow.core import Task
@@ -55,6 +56,8 @@ def get_children_right(self):
         return self.children_right
 
     def is_classifier(self) -> bool:
+        if tensorflow_decision_forests.__version__<'1.2.0':
+            return self.model._task == Task.CLASSIFICATION
         return self.model.task == Task.CLASSIFICATION
 
     def get_class_weights(self):