From 3bf9af7d27e82c50dcebe762528ae185ee52d9ea Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Wed, 3 Aug 2022 18:23:37 +0800
Subject: [PATCH] Cleanup.

---
 src/data/iterative_dmatrix.h | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/data/iterative_dmatrix.h b/src/data/iterative_dmatrix.h
index bcb98093be91..976b87d56ee8 100644
--- 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);