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
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

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)