Fix gpu_id with custom objective. (#7015)

This commit is contained in:
Jiaming Yuan
2021-06-09 14:51:17 +08:00
committed by GitHub
parent bd2ca543c4
commit 72f9daf9b6
3 changed files with 19 additions and 6 deletions

View File

@@ -228,7 +228,7 @@ void GBTree::DoBoost(DMatrix* p_fmat,
// break a lots of existing code.
auto device = tparam_.tree_method != TreeMethod::kGPUHist
? GenericParameter::kCpuId
: in_gpair->DeviceIdx();
: generic_param_->gpu_id;
auto out = MatrixView<float>(
&predt->predictions,
{static_cast<size_t>(p_fmat->Info().num_row_), static_cast<size_t>(ngroup)}, device);