diff --git a/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/params/NonParamVariables.scala b/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/params/NonParamVariables.scala index 791c65172..276a938e0 100644 --- a/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/params/NonParamVariables.scala +++ b/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/params/NonParamVariables.scala @@ -21,8 +21,9 @@ import org.apache.spark.sql.DataFrame trait NonParamVariables { protected var evalSetsMap: Map[String, DataFrame] = Map.empty - def setEvalSets(evalSets: Map[String, DataFrame]): Unit = { + def setEvalSets(evalSets: Map[String, DataFrame]): this.type = { evalSetsMap = evalSets + this } def getEvalSets(params: Map[String, Any]): Map[String, DataFrame] = {