set_b on gpu

src/loss_function/LossFunctionPearsonRMSE.cpp

@@ -341,7 +341,8 @@ Kokkos::View<double*> LossFunctionPearsonRMSE::operator()(
         start += _task_sizes_train[task_idx];
     }
 
-    ::get_mean_squared_difference(_batched_scores, _training_properties, _estimated_training_properties);
+    ::get_mean_squared_difference(
+        _batched_scores, _training_properties, _estimated_training_properties);
 
     return _batched_scores;
 }
@@ -427,7 +428,6 @@ void LossFunctionPearsonRMSE::set_a(const std::vector<std::vector<int>>& feature
                                                                   models(feature_idx, model_idx));
     };
     Kokkos::parallel_for("LossFunctionPearsonRMSE::set_a", policy, kernel);
-    Kokkos::fence();
 }
 
 void LossFunctionPearsonRMSE::set_a(const std::vector<model_node_ptr>& feats,
@@ -442,10 +442,15 @@ void LossFunctionPearsonRMSE::set_a(const std::vector<model_node_ptr>& feats,
 
 void LossFunctionPearsonRMSE::set_b(int taskind, int start)
 {
-    for (size_t batch_idx = 0; batch_idx < MAX_BATCHES; ++batch_idx)
+    auto b = _b;
+    auto training_properties = _training_properties;
+    auto policy = Kokkos::MDRangePolicy<Kokkos::Rank<2>>({0, 0},
+                                                         {_task_sizes_train[taskind], MAX_BATCHES});
+    auto kernel = KOKKOS_LAMBDA(const int material_idx, const int batch_idx)
     {
-        std::copy_n(&_prop_train[start], _task_sizes_train[taskind], &_b(0, batch_idx));
-    }
+        b(material_idx, batch_idx) = training_properties(start + material_idx);
+    };
+    Kokkos::parallel_for("LossFunctionPearsonRMSE::set_b", policy, kernel);
 }
 
 int LossFunctionPearsonRMSE::least_squares(int taskind, int start)