From 1b657a55138f60dbf46ec9e5c80c05aae3cf2c4d Mon Sep 17 00:00:00 2001 From: Bobby Wang Date: Mon, 24 Jul 2023 18:32:25 +0800 Subject: [PATCH] [jvm-packages] set device to cuda when tree method is "gpu_hist" (#9412) --- .../scala/ml/dmlc/xgboost4j/scala/spark/XGBoost.scala | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoost.scala b/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoost.scala index 5fc16ec09..f514eaa68 100644 --- a/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoost.scala +++ b/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoost.scala @@ -180,10 +180,12 @@ private[this] class XGBoostExecutionParamsFactory(rawParams: Map[String, Any], s " as 'hist', 'approx', 'gpu_hist', and 'auto'") treeMethod = Some(overridedParams("tree_method").asInstanceOf[String]) } - val device: Option[String] = overridedParams.get("device") match { - case None => None - case Some(dev: String) => if (treeMethod == "gpu_hist") Some("cuda") else Some(dev) - } + + // back-compatible with "gpu_hist" + val device: Option[String] = if (treeMethod.exists(_ == "gpu_hist")) { + Some("cuda") + } else overridedParams.get("device").map(_.toString) + if (overridedParams.contains("train_test_ratio")) { logger.warn("train_test_ratio is deprecated since XGBoost 0.82, we recommend to explicitly" + " pass a training and multiple evaluation datasets by passing 'eval_sets' and " +