Cleanup.

dmlc · Aug 8, 2022 · 3bf9af7 · 3bf9af7
1 parent 837aabd
commit 3bf9af7
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/data/iterative_dmatrix.h b/src/data/iterative_dmatrix.h
@@ -97,8 +97,9 @@ class IterativeDMatrix : public DMatrix {
     batch_param_ = BatchParam{d, max_bin};
     batch_param_.sparse_thresh = 0.2;  // default from TrainParam
 
-    ctx_.UpdateAllowUnknown(Args{{"nthread", std::to_string(nthread)}});
-    if (d == Context::kCpuId) {
+    ctx_.UpdateAllowUnknown(
+        Args{{"nthread", std::to_string(nthread)}, {"gpu_id", std::to_string(d)}});
+    if (ctx_.IsCPU()) {
       this->InitFromCPU(iter_handle, missing, ref);
     } else {
       this->InitFromCUDA(iter_handle, missing, ref);