fix regr

mlr-org · Oct 20, 2021 · 650d05d · 650d05d
1 parent 54662a5
commit 650d05d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/R/learner_lightgbm_regr_lightgbm.R b/R/learner_lightgbm_regr_lightgbm.R
@@ -231,8 +231,8 @@ LearnerRegrLightGBM = R6Class("LearnerRegrLightGBM",
         train_ids = setdiff(seq(task$nrow), task$row_roles$validation)
 
         dtrain = lightgbm::lgb.Dataset(
-          data = as.matrix(task$data(cols = task$feature_names)),
-          label = train_label,
+          data = as.matrix(task$data(rows = train_ids, cols = task$feature_names)),
+          label = as.matrix(task$data(rows = train_ids, cols = task$target_names)),
           free_raw_data = FALSE,
           categorical_feature = pars$categorical_feature
         )