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/src/c_api/c_api_utils.h
+++ b/src/c_api/c_api_utils.h
@@ -96,6 +96,10 @@ inline void CalcPredictShape(bool strict_shape, PredictionType type, size_t rows
      forest = std::max(static_cast<decltype(forest)>(1), forest);
      shape[3] = forest;
      *out_dim = shape.size();
+    } else if (chunksize == 1) {
+      *out_dim = 1;
+      shape.resize(*out_dim);
+      shape.front() = rows;
    } else {
      *out_dim = 2;
      shape.resize(*out_dim);