Explicitly check for list type

dmlc · Feb 16, 2024 · 7333eb9 · 7333eb9
1 parent e3c0263
commit 7333eb9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/treelite/sklearn/importer.py b/python/treelite/sklearn/importer.py
@@ -139,9 +139,9 @@ def import_model(sklearn_model):
         )
     elif isinstance(sklearn_model, (RandomForestC, ExtraTreesC)):
         # pylint: disable=C3001
-        try:  # multi-target
+        if isinstance(sklearn_model.n_classes_, list):  # multi-target
             max_n_classes = max(sklearn_model.n_classes_)
-        except TypeError:  # single-target
+        else:  # single-target
             max_n_classes = sklearn_model.n_classes_
         leaf_value_expected_shape = lambda node_count: (  # noqa: E731
             node_count,