Commit ea37b10b authored by Zeeshan Siddiqui's avatar Zeeshan Siddiqui Committed by Guolin Ke
Browse files

Prevent memory corruption when multi-class dataset is passed for binary light GBM (#444)

parent 8bf1297a
...@@ -170,7 +170,10 @@ void GBDT::ResetTrainingData(const BoostingConfig* config, const Dataset* train_ ...@@ -170,7 +170,10 @@ void GBDT::ResetTrainingData(const BoostingConfig* config, const Dataset* train_
std::vector<data_size_t> cnt_per_class(num_tree_per_iteration_ + 1, 0); std::vector<data_size_t> cnt_per_class(num_tree_per_iteration_ + 1, 0);
auto label = train_data_->metadata().label(); auto label = train_data_->metadata().label();
for (int i = 0; i < num_data_; ++i) { for (int i = 0; i < num_data_; ++i) {
++cnt_per_class[static_cast<int>(label[i])]; int index = static_cast<int>(label[i]);
//Check if user gave a multi-class dataset for a binary class problem.
if (index <= num_tree_per_iteration_)
++cnt_per_class[static_cast<int>(label[i])];
} }
if (num_tree_per_iteration_ > 1) { if (num_tree_per_iteration_ > 1) {
for (int i = 0; i < num_tree_per_iteration_; ++i) { for (int i = 0; i < num_tree_per_iteration_; ++i) {
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment