From c8c472f39aa08d5b8f180076a429f09317451cc7 Mon Sep 17 00:00:00 2001
From: Rory Mitchell <r.a.mitchell.nz@gmail.com>
Date: Wed, 20 Feb 2019 09:23:15 +1300
Subject: [PATCH] Fix incorrect device in multi-GPU algorithm (#4161)

---
 src/tree/updater_gpu_hist.cu | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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