Update docstrings

zillow · Sep 6, 2024 · 0bdc756 · 0bdc756
1 parent 2786931
commit 0bdc756
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/quantile_forest/_quantile_forest.py b/quantile_forest/_quantile_forest.py
@@ -686,7 +686,8 @@ def predict(
 
         Returns
         -------
-        y_pred : array of shape (n_samples, n_outputs, n_quantiles)
+        y_pred : array of shape (n_train, n_quantiles) or \
+                (n_train, n_outputs, n_quantiles)
             If quantiles is set to 'mean', then return ``E(Y | X)``. Else, for
             all quantiles, return ``y`` at ``q`` for which ``F(Y=y|x) = q``,
             where ``q`` is the quantile.