diff --git a/src/Chemistry/SurrogateModels/AI_functions.cpp b/src/Chemistry/SurrogateModels/AI_functions.cpp
index 1d67baa2c..c966c1a4b 100644
--- a/src/Chemistry/SurrogateModels/AI_functions.cpp
+++ b/src/Chemistry/SurrogateModels/AI_functions.cpp
@@ -491,7 +491,7 @@ void parallel_training(EigenModel* Eigen_model, EigenModel* Eigen_model_reactive
     // Initialize training data input and targets
     std::vector<std::vector<double>> inputs(training_data_buffer->x.size(),
                                             std::vector<double>(params.training_data_size));
-    std::vector<std::vector<double>> targets(training_data_buffer->y.size(),
+    std::vector<std::vector<double>> targets(training_data_buffer->x.size(),
                                              std::vector<double>(params.training_data_size));
 
     int buffer_size = training_data_buffer->x[0].size();
diff --git a/src/Chemistry/WorkerFunctions.cpp b/src/Chemistry/WorkerFunctions.cpp
index c47ae8c5f..186812f93 100644
--- a/src/Chemistry/WorkerFunctions.cpp
+++ b/src/Chemistry/WorkerFunctions.cpp
@@ -170,6 +170,10 @@ void poet::ChemistryModule::WorkerDoWork(MPI_Status &probe_status,
     }
   }
 
+<<<<<<< HEAD
+=======
+  // std::cout << this->comm_rank << ":" << counter++ << std::endl;
+>>>>>>> segfault_find_reason
   if (dht_enabled || interp_enabled) {
     dht->prepareKeys(s_curr_wp.input, dt);
   }
@@ -186,6 +190,10 @@ void poet::ChemistryModule::WorkerDoWork(MPI_Status &probe_status,
     interp->tryInterpolation(s_curr_wp);
   }
 
+<<<<<<< HEAD
+=======
+
+>>>>>>> segfault_find_reason
   phreeqc_time_start = MPI_Wtime();
 
   WorkerRunWorkPackage(s_curr_wp, current_sim_time, dt);
diff --git a/src/poet.cpp b/src/poet.cpp
index 83916cf80..64b063b80 100644
--- a/src/poet.cpp
+++ b/src/poet.cpp
@@ -455,13 +455,16 @@ static Rcpp::List RunMasterLoop(RInsidePOET &R, const RuntimeParameters &params,
 
     /* AI surrogate iterative training*/
     if (params.use_ai_surrogate && !params.disable_training) {
+      // Add values for which the predictions were invalid
+      // to training data buffer      
       MSG("AI: Add to training data buffer");
-      if (!params.train_only_invalid) {
-        // Use all values if not specified otherwise 
-        R.parseEval("validity_vector <- rep(0, length(validity_vector))");
-      }
       std::vector<std::vector<double>> invalid_x = 
         R.parseEval("get_invalid_values(predictors_scaled, validity_vector)");
+      
+      if (!params.train_only_invalid) {
+        // Use all values if not specified otherwise
+        R.parseEval("validity_vector[] <- 0");
+      }
 
       R.parseEval("target_scaled <- preprocess(state_C[ai_surrogate_species])");
       std::vector<std::vector<double>> invalid_y = 
@@ -682,6 +685,8 @@ int main(int argc, char *argv[]) {
         R["ai_surrogate_species"] =
             init_list.getChemistryInit().dht_species.getNames();
 
+// TODO REMOVE!!
+R.parseEval("ai_surrogate_species <- ai_surrogate_species[ai_surrogate_species != \"Charge\"]");
         const std::string ai_surrogate_input_script =
             init_list.getChemistryInit().ai_surrogate_input_script;
 
@@ -710,6 +715,9 @@ int main(int argc, char *argv[]) {
         if (Rcpp::as<bool>(R.parseEval("exists(\"disable_training\")"))) {
           run_params.disable_training = R["disable_training"];
         }
+        if (Rcpp::as<bool>(R.parseEval("exists(\"train_only_invalid\")"))) {
+          run_params.train_only_invalid = R["train_only_invalid"];
+        }
         if (Rcpp::as<bool>(R.parseEval("exists(\"save_model_path\")"))) {
           run_params.save_model_path = Rcpp::as<std::string>(R["save_model_path"]);
           MSG("AI: Model will be saved as \"" + run_params.save_model_path + "\"");
@@ -723,7 +731,7 @@ int main(int argc, char *argv[]) {
         }
         if (!Rcpp::as<bool>(R.parseEval("exists(\"model_reactive_file_path\")"))) {
           R.parseEval("model_reactive_file_path <- model_file_path");
-        }
+        }        
         MSG("AI: Initialize Python for the AI surrogate functions");
         std::string python_keras_file = std::string(SRC_DIR) +
           "/src/Chemistry/SurrogateModels/AI_Python_functions/keras_AI_surrogate.py";
diff --git a/src/poet.hpp.in b/src/poet.hpp.in
index 9d343dccf..35cb1b4da 100644
--- a/src/poet.hpp.in
+++ b/src/poet.hpp.in
@@ -71,12 +71,21 @@ struct RuntimeParameters {
 
   /*AI surriogate configuration*/
   bool use_ai_surrogate = false; // Can be set with command line flag ---ai-surrogate
+<<<<<<< HEAD
   bool disable_training = false; // Can be set in the R input script
   bool use_clustering = false; // Can be set in the R input script
   bool use_Keras_predictions = false; // Can be set in the R input script
   bool train_only_invalid = false; // Can be set in the R input script
   int batch_size = 2560; // default value determined in test on the UP Turing cluster
   int training_epochs = 20;; // Can be set in the R input script
+=======
+  bool disable_training; // Can be set in the R input script
+  bool use_k_means_clustering; // Can be set in the R input script
+  bool use_Keras_predictions; // Can be set in the R input script
+  bool train_only_invalid; // Can be set in the R input script
+  int batch_size; // Can be set in the R input script
+  int training_epochs; // Can be set in the R input script
+>>>>>>> segfault_find_reason
   int training_data_size; // Can be set in the R input script
   std::string save_model_path = ""; // Can be set in the R input script
   std::string cuda_src_dir = CUDA_SRC_DIR; // From CMake