Fix test.

dmlc · Oct 29, 2021 · b10f886 · b10f886
1 parent 59c45f2
commit b10f886
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/tests/python/test_dmatrix.py b/tests/python/test_dmatrix.py
@@ -17,9 +17,12 @@
 
 def set_base_margin_info(DType, DMatrixT, tm: str):
     rng = np.random.default_rng()
-    X = rng.normal(0, 1.0, size=100).reshape(10, 10)
-    y = X[:, 0]
-    base_margin = DType(X)
+    X = DType(rng.normal(0, 1.0, size=100).reshape(50, 2))
+    if hasattr(X, "iloc"):
+        y = X.iloc[:, 0]
+    else:
+        y = X[:, 0]
+    base_margin = X
     # no error at set
     Xy = DMatrixT(X, y, base_margin=base_margin)
     # Error at train, caused by check in predictor.