diff --git a/src/models/select_targets.py b/src/models/select_targets.py
index 69e70570..89f40f03 100644
--- a/src/models/select_targets.py
+++ b/src/models/select_targets.py
@@ -12,9 +12,10 @@ target_variable_name = esm_names.str.contains(snakemake.params["target_variable"
 if all(~target_variable_name):
     raise ValueError("The requested target (", snakemake.params["target_variable"],  ")cannot be found in the dataset.",
                      "Please check the names of phone_esm_ columns in all_sensor_features_cleaned_rapids.csv")
-esm_names = esm_names[~target_variable_name]
-# We will only keep one column related to phone_esm and that will be our target variable.
 
 model_input = cleaned_sensor_features.drop(esm_names, axis=1)
+model_input["target"] = cleaned_sensor_features[esm_names[target_variable_name]]
+# We will only keep one column related to phone_esm and that will be our target variable.
+# Add it back to the very and of the data frame and rename it to target.
 
 model_input.to_csv(snakemake.output[0], index=False)