double nrows in test case when nclasses is larger than 5

rapidsai · Mar 16, 2024 · 57989ef · 57989ef
1 parent 39048c8
commit 57989ef
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/cuml/tests/dask/test_dask_logistic_regression.py b/python/cuml/tests/dask/test_dask_logistic_regression.py
@@ -720,9 +720,9 @@ def test_standardization_on_scaled_dataset(
     penalty = regularization[0]
     C = regularization[1]
     l1_ratio = regularization[2]
-    nrows = int(1e5)
-    ncols = ncol_and_nclasses[0]
     n_classes = ncol_and_nclasses[1]
+    nrows = int(1e5) if n_classes < 5 else int(2e5)
+    ncols = ncol_and_nclasses[0]
     n_info = ncols
     n_redundant = 0
     n_parts = 2