From e4a273e1daec4e49259337c3bd3ff999e7f49451 Mon Sep 17 00:00:00 2001 From: Jiaming Yuan Date: Wed, 29 Jul 2020 19:33:29 +0800 Subject: [PATCH] Fix evaluate root split. (#5948) --- src/tree/gpu_hist/evaluate_splits.cu | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/tree/gpu_hist/evaluate_splits.cu b/src/tree/gpu_hist/evaluate_splits.cu index 63079cde8..7a843c2bf 100644 --- a/src/tree/gpu_hist/evaluate_splits.cu +++ b/src/tree/gpu_hist/evaluate_splits.cu @@ -228,13 +228,14 @@ void EvaluateSplits(common::Span out_splits, return 0; }); size_t temp_storage_bytes = 0; + auto num_segments = out_splits.size(); cub::DeviceSegmentedReduce::Sum(nullptr, temp_storage_bytes, feature_best_splits.data(), out_splits.data(), - 2, reduce_offset, reduce_offset + 1); + num_segments, reduce_offset, reduce_offset + 1); dh::TemporaryArray temp(temp_storage_bytes); cub::DeviceSegmentedReduce::Sum(temp.data().get(), temp_storage_bytes, feature_best_splits.data(), out_splits.data(), - 2, reduce_offset, reduce_offset + 1); + num_segments, reduce_offset, reduce_offset + 1); } template