Fix quantile regression demo. (#9991)

This commit is contained in:
Jiaming Yuan
2024-01-17 13:19:08 +08:00
committed by GitHub
parent cacb4b1fdd
commit d07e8b503e

View File

@@ -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)