From 4c627cb98822d77c2ec7b6497b27f91511bf2a5d Mon Sep 17 00:00:00 2001
From: RektPunk <rektpunk@gmail.com>
Date: Sat, 14 Dec 2024 16:42:51 +0900
Subject: [PATCH] add pred_leaf to reshape if condition

---
 python-package/lightgbm/basic.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python-package/lightgbm/basic.py b/python-package/lightgbm/basic.py
index 1db55385af1b..2f6ead08909d 100644
--- a/python-package/lightgbm/basic.py
+++ b/python-package/lightgbm/basic.py
@@ -1248,7 +1248,7 @@ def predict(
         if pred_leaf:
             preds = preds.astype(np.int32)
         is_sparse = isinstance(preds, (list, scipy.sparse.spmatrix))
-        if not is_sparse and preds.size != nrow:
+        if not is_sparse and preds.size != nrow or pred_leaf:
             if preds.size % nrow == 0:
                 preds = preds.reshape(nrow, -1)
             else: