diff --git a/src/data/device_adapter.cuh b/src/data/device_adapter.cuh
index ed9e3548fd9f..10ae6ba3da87 100644
--- a/src/data/device_adapter.cuh
+++ b/src/data/device_adapter.cuh
@@ -206,15 +206,13 @@ size_t GetRowCounts(const AdapterBatchT batch, common::Span<size_t> offset,
   IsValidFunctor is_valid(missing);
   // Count elements per row
   dh::LaunchN(device_idx, batch.Size(), [=] __device__(size_t idx) {
-      assert(batch.NumRows() != 0);
-      assert(batch.NumCols() != 0);
-      auto element = batch.GetElement(idx);
-      if (is_valid(element)) {
-        atomicAdd(reinterpret_cast<unsigned long long*>(  // NOLINT
-            &offset[element.row_idx]),
-                  static_cast<unsigned long long>(1));  // NOLINT
-      }
-    });
+    auto element = batch.GetElement(idx);
+    if (is_valid(element)) {
+      atomicAdd(reinterpret_cast<unsigned long long*>(  // NOLINT
+                    &offset[element.row_idx]),
+                static_cast<unsigned long long>(1));  // NOLINT
+    }
+  });
   dh::XGBCachingDeviceAllocator<char> alloc;
   size_t row_stride = thrust::reduce(
       thrust::cuda::par(alloc), thrust::device_pointer_cast(offset.data()),
diff --git a/src/data/iterative_device_dmatrix.cu b/src/data/iterative_device_dmatrix.cu
index c37df0102237..65ab90e6105d 100644
--- a/src/data/iterative_device_dmatrix.cu
+++ b/src/data/iterative_device_dmatrix.cu
@@ -64,7 +64,6 @@ void IterativeDeviceDMatrix::Initialize(DataIterHandle iter_handle, float missin
   size_t batches = 0;
   size_t accumulated_rows = 0;
   bst_feature_t cols = 0;
-
   while (iter.Next()) {
     auto device = proxy->DeviceIdx();
     dh::safe_cuda(cudaSetDevice(device));
diff --git a/tests/cpp/helpers.cc b/tests/cpp/helpers.cc
index edb94a5baad1..2274e57e7307 100644
--- a/tests/cpp/helpers.cc
+++ b/tests/cpp/helpers.cc
@@ -249,6 +249,7 @@ RandomDataGenerator::GenerateArrayInterfaceBatch(
   size_t const remaining = rows_ - offset / cols_;
   CHECK_LE(offset, rows_ * cols_);
   objects.emplace_back(make_interface(offset, remaining));
+
   for (size_t i = 0; i < batches; ++i) {
     Json::Dump(objects[i], &result[i]);
   }
diff --git a/tests/python-gpu/test_from_cudf.py b/tests/python-gpu/test_from_cudf.py
index be68df89a4f1..0ba5931c9b56 100644
--- a/tests/python-gpu/test_from_cudf.py
+++ b/tests/python-gpu/test_from_cudf.py
@@ -234,13 +234,17 @@ def test_from_cudf_iter():
     it = IterForDMatrixTest()
 
     # Use iterator
-    m = xgb.DeviceQuantileDMatrix(it)
-    reg_with_it = xgb.train({'tree_method': 'gpu_hist'}, m,
+    m_it = xgb.DeviceQuantileDMatrix(it)
+    reg_with_it = xgb.train({'tree_method': 'gpu_hist'}, m_it,
                             num_boost_round=rounds)
-    predict_with_it = reg_with_it.predict(m)
+    predict_with_it = reg_with_it.predict(m_it)
 
     # Without using iterator
     m = xgb.DMatrix(it.as_array(), it.as_array_labels())
+
+    assert m_it.num_col() == m.num_col()
+    assert m_it.num_row() == m.num_row()
+
     reg = xgb.train({'tree_method': 'gpu_hist'}, m,
                     num_boost_round=rounds)
     predict = reg.predict(m)