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.