Fix inference for xgboost fitted with early-stopping (#1350)

intel · Jul 12, 2023 · 27d3f39 · 27d3f39
1 parent 1a6a655
commit 27d3f39
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/gbt_convertors.pyx b/src/gbt_convertors.pyx
@@ -152,7 +152,8 @@ def get_gbt_model_from_xgboost(booster: Any) -> Any:
     else:
         is_regression = True
 
-    n_iterations = int(len(trees_arr) / (n_classes if n_classes > 2 else 1))
+    n_iterations = booster.best_iteration + 1
+    trees_arr = trees_arr[: n_iterations * (n_classes if n_classes > 2 else 1)]
 
     # Create + base iteration
     if is_regression: