[Bug]: store_predictions=TRUE in fit() does not save results with more than one treatment

[sharpant]

Describe the bug

I am working with DoubleMLPLIV. When there is more than one treatment and more than one instrument, and I use store_predictions=TRUE in fit(), the results for ml_m do not get saved and show NAs. Can this please be fixed?

Minimum reproducible code snippet

library('DoubleML')
library('mlr3')
library('mlr3learners')

# Set the simulation parameters
N = 25  # number of observations (first dimension)
M = 25  # number of observations (second dimension)
dim_X = 100  # dimension of X
set.seed(3141) # set seed

obj_dml_data = make_pliv_multiway_cluster_CKMS2021(N, M, dim_X)

data = as.data.frame(obj_dml_data$data)

data$D2 <- data$D*data$D
data$Z2 <- data$Z*data$Z

dml_1 = double_ml_data_from_data_frame(df=data,x_cols=paste0("X",seq(1,100)), y_col="Y", d_cols=c("D","D2"),z_cols=c("Z","Z2"),cluster_cols = c("cluster_var_i","cluster_var_j")) 

# The simulated data is of type DoubleMLClusterData
print(dml_1)

# Set machine learning methods for l, m & r
ml_l = lrn("regr.cv_glmnet", nfolds = 10, s = "lambda.min")
ml_m = lrn("regr.cv_glmnet", nfolds = 10, s = "lambda.min")
ml_r = lrn("regr.cv_glmnet", nfolds = 10, s = "lambda.min")

# initialize the DoubleMLPLIV object
dml_pliv_obj = DoubleMLPLIV$new(dml_1,
                                ml_l, ml_m, ml_r,
                                n_folds=3)

print(dml_pliv_obj)

dml_pliv_obj$fit(store_predictions=TRUE)

#dml_pliv_obj$summary()

preds=dml_pliv_obj$predictions

#Results are NA
preds$ml_m_Z
preds$ml_m_Z2

Expected Result

There should not be NAs saved for ml_m variables

preds$ml_m_Z
preds$ml_m_Z2

Actual Result

I get NAs saved for

preds$ml_m_Z
preds$ml_m_Z2

Versions

R version 4.2.2 (2022-10-31 ucrt)
Platform: x86_64-w64-mingw32/x64 (64-bit)
Running under: Windows 10 x64 (build 19045)

Matrix products: default

locale:
[1] LC_COLLATE=English_Canada.utf8 LC_CTYPE=English_Canada.utf8
[3] LC_MONETARY=English_Canada.utf8 LC_NUMERIC=C
[5] LC_TIME=English_Canada.utf8

attached base packages:
[1] stats graphics grDevices utils datasets methods base

other attached packages:
[1] mlr3learners_0.5.4 mlr3_0.14.1 DoubleML_0.5.3

loaded via a namespace (and not attached):
[1] Rcpp_1.0.9 parallelly_1.32.1 splines_4.2.2
[4] MASS_7.3-58.1 uuid_1.1-0 lattice_0.20-45
[7] R6_2.5.1 foreach_1.5.2 lgr_0.4.4
[10] globals_0.16.1 tools_4.2.2 palmerpenguins_0.1.1
[13] grid_4.2.2 glmnet_4.1-4 parallel_4.2.2
[16] checkmate_2.1.0 data.table_1.14.2 clusterGeneration_1.3.7
[19] mlr3tuning_0.15.0 iterators_1.0.14 survival_3.4-0
[22] digest_0.6.29 crayon_1.5.2 Matrix_1.5-1
[25] paradox_0.10.0 mlr3measures_0.5.0 codetools_0.2-18
[28] shape_1.4.6 future.apply_1.9.1 compiler_4.2.2
[31] readstata13_0.10.0 bbotk_0.6.0 backports_1.4.1
[34] mlr3misc_0.11.0 future_1.28.0 mvtnorm_1.1-3
[37] listenv_0.8.0

packageVersion('DoubleML')
[1] ‘0.5.3’
packageVersion('mlr3')
[1] ‘0.14.1’