Fix test.

dmlc · Jul 20, 2021 · 8629f51 · 8629f51
1 parent 9691d15
commit 8629f51
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tests/python/test_data_iterator.py b/tests/python/test_data_iterator.py
@@ -61,7 +61,9 @@ def run_data_iterator(
 ) -> None:
     n_rounds = 2
 
-    it = IteratorForTest(*make_batches(n_samples_per_batch, n_features, n_batches))
+    it = IteratorForTest(
+        *make_batches(n_samples_per_batch, n_features, n_batches, cupy)
+    )
     if n_batches == 0:
         with pytest.raises(ValueError, match="1 batch"):
             Xy = xgb.DMatrix(it)