hotfix: dmlc#6755

h2oai · Mar 16, 2021 · 9ddd292 · 9ddd292
1 parent 19e81bc
commit 9ddd292
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
@@ -169,6 +169,8 @@ struct GPUHistMakerDevice {
   std::unique_ptr<RowPartitioner> row_partitioner;
   DeviceHistogram<GradientSumT> hist{};
 
+  # https://github.com/dmlc/xgboost/pull/6755/files
+  dh::caching_device_vector<GradientPair> d_gpair; // storage for gpair;
   common::Span<GradientPair> gpair;
 
   dh::caching_device_vector<int> monotone_constraints;
@@ -269,7 +271,15 @@ struct GPUHistMakerDevice {
     std::fill(node_sum_gradients.begin(), node_sum_gradients.end(),
               GradientPair());
 
-    auto sample = sampler->Sample(dh_gpair->DeviceSpan(), dmat);
+    # https://github.com/dmlc/xgboost/pull/6755/files
+    if (d_gpair.size() != dh_gpair->Size()) {
+      d_gpair.resize(dh_gpair->Size());
+    }
+    thrust::copy(thrust::device, dh_gpair->ConstDevicePointer(),
+                 dh_gpair->ConstDevicePointer() + dh_gpair->Size(),
+                 d_gpair.begin());
+    auto sample = sampler->Sample(dh::ToSpan(d_gpair), dmat);
+
     page = sample.page;
     gpair = sample.gpair;