fix parrt#304

MohammedAlsayed · Sep 13, 2023 · 892a4a1 · 892a4a1
1 parent be0cdfc
commit 892a4a1
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/dtreeviz/models/shadow_decision_tree.py b/dtreeviz/models/shadow_decision_tree.py
@@ -523,7 +523,8 @@ def n_sample_classes(self):
             return node_values
         node_y_data = self.shadow_tree.y_train[samples]
         unique, counts = np.unique(node_y_data, return_counts=True)
-
+        if unique.dtype != 'int':
+            raise ValueError(f"Target values should be type int got {unique.dtype}")
         for i in range(len(unique)):
             node_values[unique[i]] = counts[i]