From 27d3f39ce95584589f7c8f838b81b4b741d59a61 Mon Sep 17 00:00:00 2001
From: avolkov-intel <117643568+avolkov-intel@users.noreply.github.com>
Date: Wed, 12 Jul 2023 09:21:15 +0200
Subject: [PATCH] Fix inference for xgboost fitted with early-stopping (#1350)

---
 src/gbt_convertors.pyx | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/gbt_convertors.pyx b/src/gbt_convertors.pyx
index 3fdbccdc6d..f6c10b96bf 100755
--- 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: