Fix evaluate root split. (#5948)
This commit is contained in:
parent
071e10c1d1
commit
e4a273e1da
@ -228,13 +228,14 @@ void EvaluateSplits(common::Span<DeviceSplitCandidate> 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<int8_t> 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 <typename GradientSumT>
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user