Relax test for decision stump in distributed environment. (#6919)

dmlc · Apr 30, 2021 · a1d23f6 · a1d23f6
1 parent 45ddc39
commit a1d23f6
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/tests/python/test_with_dask.py b/tests/python/test_with_dask.py
@@ -1023,7 +1023,17 @@ def run_updater_test(
                                  evals=[(m, 'train')])['history']
         note(history)
         history = history['train'][dataset.metric]
-        assert tm.non_increasing(history)
+
+        def is_stump():
+            return params["max_depth"] == 1 or params["max_leaves"] == 1
+
+        def minimum_bin():
+            return "max_bin" in params and params["max_bin"] == 2
+
+        if minimum_bin() and is_stump():
+            assert tm.non_increasing(history, tolerance=1e-3)
+        else:
+            assert tm.non_increasing(history)
         # Make sure that it's decreasing
         assert history[-1] < history[0]