revert back api (intel-analytics#2943)

dding3 · Oct 24, 2019 · 473f170 · 473f170
1 parent 368998a
commit 473f170
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/dl/src/main/scala/com/intel/analytics/bigdl/dllib/nn/abstractnn/AbstractModule.scala b/dl/src/main/scala/com/intel/analytics/bigdl/dllib/nn/abstractnn/AbstractModule.scala
@@ -647,7 +647,7 @@ abstract class AbstractModule[A <: Activity: ClassTag, B <: Activity: ClassTag,
    *                  if -1, default is 4 * partitionNumber of dataset
    */
 
-  final def predictClass(dataset: RDD[Sample[T]], batchSize: Int = -1): RDD[Sample[T]] = {
+  final def predictClass(dataset: RDD[Sample[T]], batchSize: Int = -1): RDD[Int] = {
     Predictor(this).predictClass(dataset, batchSize)
   }