fix: store learner in state of PipeOpBreslow

mlr-org · Jan 18, 2024 · 6987cc7 · 6987cc7
1 parent 8565481
commit 6987cc7
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/R/PipeOpBreslow.R b/R/PipeOpBreslow.R
@@ -124,6 +124,7 @@ PipeOpBreslow = R6Class("PipeOpBreslow",
 
       # keep the training data that Breslow estimator needs
       self$state = list(
+        learner = learner,
         times = task$times(),
         status = task$status(),
         lp_train = p$lp
@@ -134,7 +135,7 @@ PipeOpBreslow = R6Class("PipeOpBreslow",
 
     .predict = function(inputs) {
       task = inputs[[1]]
-      learner = private$.learner
+      learner = self$state$learner
 
       if (is.null(learner$model)) {
         stopf("Cannot predict, Learner '%s' has not been trained yet", learner$id)