[R] Fix softprob reshape. (#7126)

2021-07-27 15:25:17 +08:00
parent 7ee7a95b84
commit 48d5de80a2
2 changed files with 16 additions and 1 deletions
--- a/R-package/R/xgb.Booster.R
+++ b/R-package/R/xgb.Booster.R
@@ -428,7 +428,7 @@ predict.xgb.Booster <- function(object, newdata, missing = NA, outputmargin = FA
    } else if (!reshape && n_groups != 1) {
      arr <- ret
    } else if (reshape && n_groups != 1) {
-      arr <- matrix(arr, ncol = 3, byrow = TRUE)
+      arr <- matrix(arr, ncol = n_groups, byrow = TRUE)
    }
    arr <- drop(arr)
    if (length(dim(arr)) == 1) {