diff --git a/demo/guide-python/quantile_regression.py b/demo/guide-python/quantile_regression.py
index 5d186714c..a9e4532ba 100644
--- a/demo/guide-python/quantile_regression.py
+++ b/demo/guide-python/quantile_regression.py
@@ -46,10 +46,11 @@ def quantile_loss(args: argparse.Namespace) -> None:
 
     X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=rng)
     # We will be using the `hist` tree method, quantile DMatrix can be used to preserve
-    # memory.
+    # memory (which has nothing to do with quantile regression itself, see its document
+    # for details).
     # Do not use the `exact` tree method for quantile regression, otherwise the
     # performance might drop.
-    Xy = xgb.QuantileDMatrix(X, y)
+    Xy = xgb.QuantileDMatrix(X_train, y_train)
     # use Xy as a reference
     Xy_test = xgb.QuantileDMatrix(X_test, y_test, ref=Xy)