Resolve saving bug with multilabel parameter (#13842)

JohnSnowLabs · Jun 8, 2023 · 5b61123 · 5b61123
1 parent ddf39b3
commit 5b61123
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/main/scala/com/johnsnowlabs/nlp/HasClassifierActivationProperties.scala b/src/main/scala/com/johnsnowlabs/nlp/HasClassifierActivationProperties.scala
@@ -16,7 +16,7 @@
 
 package com.johnsnowlabs.nlp
 
-import org.apache.spark.ml.param.{FloatParam, Param}
+import org.apache.spark.ml.param.{BooleanParam, FloatParam, Param}
 
 trait HasClassifierActivationProperties extends ParamsAndFeaturesWritable {
 
@@ -47,7 +47,7 @@ trait HasClassifierActivationProperties extends ParamsAndFeaturesWritable {
     *
     * @group param
     */
-  val multilabel: Param[Boolean] = new Param(
+  val multilabel: BooleanParam = new BooleanParam(
     this,
     "multilabel",
     "Whether or not the result should be multi-class (the sum of all probabilities is 1.0) or multi-label (each label has a probability between 0.0 to 1.0). Default is False i.e. multi-class")