Ensure predict leaf output 1-dim vector where there's only 1 tree. (#6889)

2021-04-23 15:07:48 +08:00
parent 54afa3ac7a
commit 8760ec4827
3 changed files with 11 additions and 2 deletions
--- a/tests/python/test_predict.py
+++ b/tests/python/test_predict.py
@@ -75,6 +75,11 @@ def run_predict_leaf(predictor):
    first = sliced[0, ...]

    assert np.prod(first.shape) == classes * num_parallel_tree * ntree_limit
+
+    # When there's only 1 tree, the output is a 1 dim vector
+    booster = xgb.train({"tree_method": "hist"}, num_boost_round=1, dtrain=m)
+    assert booster.predict(m, pred_leaf=True).shape == (rows, )
+
    return leaf