diff --git a/R-package/R/xgb.DMatrix.R b/R-package/R/xgb.DMatrix.R index cb975cf06..ea9cfcc47 100644 --- a/R-package/R/xgb.DMatrix.R +++ b/R-package/R/xgb.DMatrix.R @@ -301,12 +301,17 @@ slice.xgb.DMatrix <- function(object, idxset, ...) { attr_list <- attributes(object) nr <- nrow(object) - len <- sapply(attr_list, length) + len <- sapply(attr_list, NROW) ind <- which(len == nr) if (length(ind) > 0) { nms <- names(attr_list)[ind] for (i in seq_along(ind)) { - attr(ret, nms[i]) <- attr(object, nms[i])[idxset] + obj_attr <- attr(object, nms[i]) + if (NCOL(obj_attr) > 1) { + attr(ret, nms[i]) <- obj_attr[idxset,] + } else { + attr(ret, nms[i]) <- obj_attr[idxset] + } } } return(structure(ret, class = "xgb.DMatrix"))