add pred_leaf to reshape if condition

microsoft · Dec 14, 2024 · 4c627cb · 4c627cb
1 parent b33a12e
commit 4c627cb
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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: