diff --git a/r-package/grf/R/average_treatment_effect.R b/r-package/grf/R/average_treatment_effect.R
index 236b7e057..79db57120 100644
--- a/r-package/grf/R/average_treatment_effect.R
+++ b/r-package/grf/R/average_treatment_effect.R
@@ -95,7 +95,7 @@ average_treatment_effect <- function(forest,
   clusters <- if (cluster.se) {
     forest$clusters
   } else {
-    1:length(forest$Y)
+    1:length(forest$Y.orig)
   }
   observation.weight <- observation_weights(forest)