Fix incorrect device in multi-GPU algorithm (#4161)

This commit is contained in:
Rory Mitchell 2019-02-20 09:23:15 +13:00 committed by GitHub
parent 1dac5e2410
commit c8c472f39a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -300,6 +300,7 @@ struct DeviceHistogram {
void AllocateHistogram(int nidx) { void AllocateHistogram(int nidx) {
if (HistogramExists(nidx)) return; if (HistogramExists(nidx)) return;
dh::safe_cuda(cudaSetDevice(device_id_));
if (data.size() > kStopGrowingSize) { if (data.size() > kStopGrowingSize) {
// Recycle histogram memory // Recycle histogram memory
std::pair<int, size_t> old_entry = *nidx_map.begin(); std::pair<int, size_t> old_entry = *nidx_map.begin();
@ -310,7 +311,6 @@ struct DeviceHistogram {
} else { } else {
// Append new node histogram // Append new node histogram
nidx_map[nidx] = data.size(); nidx_map[nidx] = data.size();
dh::safe_cuda(cudaSetDevice(device_id_));
// x 2: Hess and Grad. // x 2: Hess and Grad.
data.resize(data.size() + (n_bins * 2)); data.resize(data.size() + (n_bins * 2));
} }