Use probabilities when testing skl classifier export

dmlc · Oct 12, 2024 · cffe894 · cffe894
1 parent 98f3ff8
commit cffe894
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/python/test_sklearn_integration.py b/tests/python/test_sklearn_integration.py
@@ -299,7 +299,7 @@ def test_skl_export_rf_classifier(dataset, max_depth, n_estimators):
     tl_model = treelite.sklearn.import_model(clf)
     clf2 = treelite.sklearn.export_model(tl_model)
     assert isinstance(clf2, RandomForestClassifier)
-    np.testing.assert_almost_equal(clf2.predict(X), clf.predict(X))
+    np.testing.assert_almost_equal(clf2.predict_proba(X), clf.predict_proba(X))
 
 
 @given(