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: