Getting TE model via getTargetEncoderModel

h2oai · Jul 17, 2019 · 0a09b5d · 0a09b5d
1 parent 2f7138d
commit 0a09b5d
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/ml/src/main/scala/org/apache/spark/ml/h2o/features/H2OTargetEncoder.scala b/ml/src/main/scala/org/apache/spark/ml/h2o/features/H2OTargetEncoder.scala
@@ -51,7 +51,8 @@ class H2OTargetEncoder(override val uid: String)
     targetEncoderParameters.setTrain(trainingFrame._key)
 
     val builder = new TargetEncoderBuilder(targetEncoderParameters)
-    builder.trainModel().get()
+    builder.trainModel().get() // Calling get() to wait until the model training is finished.
+    builder.getTargetEncoderModel()
   }
 
   override def copy(extra: ParamMap): H2OTargetEncoder = defaultCopy(extra)