diff --git a/quantile_forest/_quantile_forest.py b/quantile_forest/_quantile_forest.py
index a0d27f6..bb61966 100755
--- a/quantile_forest/_quantile_forest.py
+++ b/quantile_forest/_quantile_forest.py
@@ -359,9 +359,6 @@ def _get_y_train_leaves(self, X, y, sorter=None, sample_weight=None):
                     bootstrap_indices=bootstrap_indices[:, i],
                     leaf_indices=leaf_indices,
                     leaf_values_list=leaf_values_list,
-                    max_node_count=max_node_count,
-                    max_samples_leaf=max_samples_leaf,
-                    n_outputs=n_outputs,
                 )
 
         return y_train_leaves
diff --git a/quantile_forest/_quantile_forest_fast.pyx b/quantile_forest/_quantile_forest_fast.pyx
index 6bb9708..643edf2 100755
--- a/quantile_forest/_quantile_forest_fast.pyx
+++ b/quantile_forest/_quantile_forest_fast.pyx
@@ -562,16 +562,14 @@ cpdef map_leaf_nodes(
     cnp.ndarray[intp_t, ndim=2] bootstrap_indices,
     vector[intp_t] leaf_indices,
     vector[vector[intp_t]] leaf_values_list,
-    int max_node_count,
-    int max_samples_leaf,
-    int n_outputs,
 ) noexcept:
-    cdef intp_t n_samples, n_leaves
+    cdef intp_t n_samples, n_outputs, n_leaves
     cdef intp_t i, j, k
     cdef vector[intp_t] leaf_values
     cdef intp_t leaf_index, leaf_value, y_index
     cdef intp_t[:, :, :] y_train_leaves_view
 
+    n_outputs = bootstrap_indices.shape[1]
     n_leaves = leaf_indices.size()
 
     y_train_leaves_view = y_train_leaves  # memoryview