From 82033064005d07ae98cd3003190d675992061a61 Mon Sep 17 00:00:00 2001
From: mjmckp <mjmckp@users.noreply.github.com>
Date: Tue, 5 Sep 2023 11:34:35 +1000
Subject: [PATCH] Fix updates in random forest model using GOSS data sample
 strategy (#6017)

---
 src/boosting/rf.hpp | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/src/boosting/rf.hpp b/src/boosting/rf.hpp
index 9a87e982483e..88ece154e432 100644
--- a/src/boosting/rf.hpp
+++ b/src/boosting/rf.hpp
@@ -115,6 +115,12 @@ class RF : public GBDT {
     const data_size_t bag_data_cnt = data_sample_strategy_->bag_data_cnt();
     const std::vector<data_size_t, Common::AlignmentAllocator<data_size_t, kAlignedSize>>& bag_data_indices = data_sample_strategy_->bag_data_indices();
 
+    // GOSSStrategy->Bagging may modify value of bag_data_cnt_
+    if (is_use_subset && bag_data_cnt < num_data_) {
+      tmp_grad_.resize(num_data_);
+      tmp_hess_.resize(num_data_);
+    }
+
     CHECK_EQ(gradients, nullptr);
     CHECK_EQ(hessians, nullptr);