diff --git a/R/gcomp.R b/R/gcomp.R
index 5125dbd..315d1ab 100644
--- a/R/gcomp.R
+++ b/R/gcomp.R
@@ -62,10 +62,10 @@ estimate_sub <- function(natural, shifted, outcome, node_list, cens, risk,
     }
 
     trt_var <- names(shifted$train)[t]
-    under_shift_train <- natural$train[jt & rt, vars]
+    under_shift_train <- as.data.frame(natural$train[jt & rt, vars])
     under_shift_train[[trt_var]] <- shifted$train[jt & rt, trt_var]
 
-    under_shift_valid <- natural$valid[jv & rv, vars]
+    under_shift_valid <- as.data.frame(natural$valid[jv & rv, vars])
     under_shift_valid[[trt_var]] <- shifted$valid[jv & rv, trt_var]
 
     natural$train[jt & rt, pseudo] <- bound(SL_predict(fit, under_shift_train), 1e-05)
diff --git a/R/sdr.R b/R/sdr.R
index 05674c8..2e74a18 100644
--- a/R/sdr.R
+++ b/R/sdr.R
@@ -83,10 +83,10 @@ estimate_sdr <- function(natural, shifted, outcome, node_list, cens, risk, tau,
     }
 
     trt_var <- names(shifted$train)[t]
-    under_shift_train <- natural$train[jt & rt, vars]
+    under_shift_train <- as.data.frame(natural$train[jt & rt, vars])
     under_shift_train[[trt_var]] <- shifted$train[jt & rt, trt_var]
 
-    under_shift_valid <- natural$valid[jv & rv, vars]
+    under_shift_valid <- as.data.frame(natural$valid[jv & rv, vars])
     under_shift_valid[[trt_var]] <- shifted$valid[jv & rv, trt_var]
 
     m_natural_train[jt & rt, t] <- bound(SL_predict(fit, natural$train[jt & rt, vars]), 1e-05)
diff --git a/R/tmle.R b/R/tmle.R
index 0584830..da5191c 100644
--- a/R/tmle.R
+++ b/R/tmle.R
@@ -67,10 +67,10 @@ estimate_tmle <- function(natural, shifted, outcome, node_list, cens, risk, tau,
     }
 
     trt_var <- names(shifted$train)[t]
-    under_shift_train <- natural$train[jt & rt, vars]
+    under_shift_train <- as.data.frame(natural$train[jt & rt, vars])
     under_shift_train[[trt_var]] <- shifted$train[jt & rt, trt_var]
 
-    under_shift_valid <- natural$valid[jv & rv, vars]
+    under_shift_valid <- as.data.frame(natural$valid[jv & rv, vars])
     under_shift_valid[[trt_var]] <- shifted$valid[jv & rv, trt_var]
 
     m_natural_train[jt & rt, t] <- bound(SL_predict(fit, natural$train[jt & rt, vars]), 1e-05)