modin.

dmlc · Dec 29, 2023 · b9431ac · b9431ac
1 parent 1880325
commit b9431ac
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/python/test_with_modin.py b/tests/python/test_with_modin.py
@@ -1,9 +1,9 @@
 import numpy as np
 import pytest
-from test_dmatrix import set_base_margin_info
 
 import xgboost as xgb
 from xgboost import testing as tm
+from xgboost.testing.data import run_base_margin_info
 
 try:
     import modin.pandas as md
@@ -145,4 +145,4 @@ def test_modin_weight(self):
         np.testing.assert_array_equal(data.get_weight(), w)
 
     def test_base_margin(self):
-        set_base_margin_info(md.DataFrame, xgb.DMatrix, "hist")
+        run_base_margin_info(md.DataFrame, xgb.DMatrix, "cpu")