diff --git a/src/tree/driver.h b/src/tree/driver.h
index 0aef93ccf..a4a0dd4a6 100644
--- a/src/tree/driver.h
+++ b/src/tree/driver.h
@@ -101,7 +101,7 @@ class Driver {
 
  private:
   TrainParam param_;
-  std::size_t num_leaves_ = 1;
+  bst_node_t num_leaves_ = 1;
   std::size_t max_node_batch_size_;
   ExpandQueue queue_;
 };