Update example plots

zillow · Aug 15, 2024 · 4ad2f1e · 4ad2f1e
1 parent 600ce47
commit 4ad2f1e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/plot_quantile_extrapolation.py b/examples/plot_quantile_extrapolation.py
@@ -145,7 +145,7 @@ def _get_tree_weight_matrix(X, Y, X_eval=None, n_trees=100, rng=None, **kwargs):
             tree.fit(X[split1, :], Y[split1].flatten())
 
             # Extract tree weight matrix.
-            y_train_leaves = tree._get_y_train_leaves(X[split2, :], 1)
+            y_train_leaves = tree._get_y_train_leaves(X[split2, :], Y)
             nrows = X[split2, :].shape[0]
             matrix = np.zeros((nrows, nrows))
             for leaf in y_train_leaves[0]: