diff --git a/examples/ML+DL-Examples/Optuna-Spark/optuna-examples/sparkrapids-xgboost-duplicate.py b/examples/ML+DL-Examples/Optuna-Spark/optuna-examples/sparkrapids-xgboost-duplicate.py index 5af2e177..eebaebb8 100644 --- a/examples/ML+DL-Examples/Optuna-Spark/optuna-examples/sparkrapids-xgboost-duplicate.py +++ b/examples/ML+DL-Examples/Optuna-Spark/optuna-examples/sparkrapids-xgboost-duplicate.py @@ -154,12 +154,12 @@ def get_gpu_id(task_context: TaskContext) -> int: Xy_train_qdm = xgb.QuantileDMatrix(X_train, y_train) # Precompute Quantile DMatrix to avoid repeated quantization every trial. def objective(trial): - params = ({ + params = { "objective": "reg:squarederror", "verbosity": 0, - "tree_method": "hist", + "tree_method": "gpu_hist", "device": f"cuda:{gpu_id}", - }) + } params.update(hyperparams.to_dict(trial)) if "max_bins" in params: