Change max_depth condition

rapidsai · Mar 5, 2020 · 5930a5d · 5930a5d
1 parent e11bdc6
commit 5930a5d
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/cpp/src/decisiontree/levelalgo/levelhelper_classifier.cuh b/cpp/src/decisiontree/levelalgo/levelhelper_classifier.cuh
@@ -299,7 +299,8 @@ void leaf_eval_classification(
   int non_leaf_counter = 0;
   // XXX: This line fixes the inaccuracy in max_depth==1 tree leaf predictions
   // I still don't fully understand its role, but it seems interesting
-  bool condition_global = true;  // XXX removed: (curr_depth == max_depth);
+  bool condition_global =
+    curr_depth >= max_depth - 1;  // XXX removed: (curr_depth == max_depth);
   if (max_leaves != -1)
     condition_global = condition_global || (tree_leaf_cnt >= max_leaves);