aboutsummaryrefslogtreecommitdiff
path: root/mllib/src/main/scala/org/apache/spark/ml/r/RWrappers.scala
diff options
context:
space:
mode:
Diffstat (limited to 'mllib/src/main/scala/org/apache/spark/ml/r/RWrappers.scala')
-rw-r--r--mllib/src/main/scala/org/apache/spark/ml/r/RWrappers.scala2
1 files changed, 2 insertions, 0 deletions
diff --git a/mllib/src/main/scala/org/apache/spark/ml/r/RWrappers.scala b/mllib/src/main/scala/org/apache/spark/ml/r/RWrappers.scala
index 7f6f147532..06baedf2a2 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/r/RWrappers.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/r/RWrappers.scala
@@ -38,6 +38,8 @@ private[r] object RWrappers extends MLReader[Object] {
val className = (rMetadata \ "class").extract[String]
className match {
case "org.apache.spark.ml.r.NaiveBayesWrapper" => NaiveBayesWrapper.load(path)
+ case "org.apache.spark.ml.r.AFTSurvivalRegressionWrapper" =>
+ AFTSurvivalRegressionWrapper.load(path)
case _ =>
throw new SparkException(s"SparkR ml.load does not support load $className")
}