From 0bdc756b839c231d4051db0e24218b87a6fff06c Mon Sep 17 00:00:00 2001
From: Reid Johnson <reidj@zillowgroup.com>
Date: Fri, 6 Sep 2024 03:31:29 -0700
Subject: [PATCH] Update docstrings

---
 quantile_forest/_quantile_forest.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/quantile_forest/_quantile_forest.py b/quantile_forest/_quantile_forest.py
index 5f05e53..b745c7b 100755
--- 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.