Fix empty DMatrix.

dmlc · Sep 2, 2022 · 1af5a8a · 1af5a8a
1 parent 8348538
commit 1af5a8a
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/src/data/iterative_dmatrix.cu b/src/data/iterative_dmatrix.cu
@@ -173,10 +173,15 @@ BatchSet<EllpackPage> IterativeDMatrix::GetEllpackBatches(BatchParam const& para
   }
   if (!ellpack_ && ghist_) {
     ellpack_.reset(new EllpackPage());
+    // Evaluation QuantileDMatrix initialized from CPU data might not have the correct GPU
+    // ID.
     if (this->ctx_.IsCPU()) {
       this->ctx_.gpu_id = param.gpu_id;
     }
-    this->Info().feature_types.SetDevice(param.gpu_id);
+    if (this->ctx_.IsCPU()) {
+      this->ctx_.gpu_id = dh::CurrentDevice();
+    }
+    this->Info().feature_types.SetDevice(this->ctx_.gpu_id);
     *ellpack_->Impl() =
         EllpackPageImpl(&ctx_, *this->ghist_, this->Info().feature_types.ConstDeviceSpan());
   }