Fix updates in random forest model using GOSS data sample strategy (#…

…6017)
microsoft · Sep 5, 2023 · 8203306 · 8203306
1 parent 5ea0057
commit 8203306
Showing 1 changed file with 6 additions and 0 deletions.
diff --git 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);