From 93661a1ea3459ef6f575ced549d44c0ac2cd240a Mon Sep 17 00:00:00 2001
From: "dmitry.razdoburdin" <drazdobu@nnldaal048.inn.intel.com>
Date: Mon, 21 Feb 2022 16:59:57 +0300
Subject: [PATCH] Fixing the UB bug

---
 plugin/updater_oneapi/regression_obj_oneapi.cc | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/plugin/updater_oneapi/regression_obj_oneapi.cc b/plugin/updater_oneapi/regression_obj_oneapi.cc
index b8df6bd59d2f..9c177e332978 100755
--- a/plugin/updater_oneapi/regression_obj_oneapi.cc
+++ b/plugin/updater_oneapi/regression_obj_oneapi.cc
@@ -71,6 +71,8 @@ class RegLossObjOneAPI : public ObjFunction {
     sycl::buffer<bst_float, 1> weights_buf(is_null_weight ? NULL : info.weights_.HostPointer(),
                                                is_null_weight ? 1 : info.weights_.Size());
 
+    const size_t n_targets = std::max(info.labels.Shape(1), static_cast<size_t>(1));
+
 	sycl::buffer<int, 1> additional_input_buf(1);
 	{
 		auto additional_input_acc = additional_input_buf.get_access<sycl::access::mode::write>();
@@ -92,7 +94,7 @@ class RegLossObjOneAPI : public ObjFunction {
       cgh.parallel_for<>(sycl::range<1>(ndata), [=](sycl::id<1> pid) {
         int idx = pid[0];
         bst_float p = Loss::PredTransform(preds_acc[idx]);
-        bst_float w = is_null_weight ? 1.0f : weights_acc[idx];
+        bst_float w = is_null_weight ? 1.0f : weights_acc[idx/n_targets];
         bst_float label = labels_acc[idx];
         if (label == 1.0f) {
           w *= scale_pos_weight;
@@ -125,7 +127,6 @@ class RegLossObjOneAPI : public ObjFunction {
 
   void PredTransform(HostDeviceVector<float> *io_preds) const override {
     size_t const ndata = io_preds->Size();
-
     sycl::buffer<bst_float, 1> io_preds_buf(io_preds->HostPointer(), io_preds->Size());
 
     qu_.submit([&](sycl::handler& cgh) {