diff --git a/mllib-dal/src/main/scala/com/intel/oap/mllib/classification/RandomForestClassifierDALImpl.scala b/mllib-dal/src/main/scala/com/intel/oap/mllib/classification/RandomForestClassifierDALImpl.scala
index a37c54612..a5e86c08e 100644
--- a/mllib-dal/src/main/scala/com/intel/oap/mllib/classification/RandomForestClassifierDALImpl.scala
+++ b/mllib-dal/src/main/scala/com/intel/oap/mllib/classification/RandomForestClassifierDALImpl.scala
@@ -81,15 +81,14 @@ class RandomForestClassifierDALImpl(val uid: String,
     }.count()
     rfcTimer.record("OneCCL Init")
 
-    val results = labeledPointsTables.mapPartitionsWithIndex {
-      (rank: Int, tables: Iterator[(String, String)]) =>
+    val results = labeledPointsTables.mapPartitionsWithIndex { (rank, tables) =>
       val (feature, label) = tables.next()
       val (featureTabAddr : Long, featureRows : Long, featureColumns : Long) = {
           val parts = feature.toString.split("_")
           (parts(0).toLong, parts(1).toLong, parts(2).toLong)
         }
       val (labelTabAddr : Long, labelRows : Long, labelColumns : Long) = {
-          val parts = feature.toString.split("_")
+          val parts = label.toString.split("_")
           (parts(0).toLong, parts(1).toLong, parts(2).toLong)
         }
       val gpuIndices = if (useDevice == "GPU") {
diff --git a/mllib-dal/src/main/scala/com/intel/oap/mllib/regression/LinearRegressionDALImpl.scala b/mllib-dal/src/main/scala/com/intel/oap/mllib/regression/LinearRegressionDALImpl.scala
index fd74aef24..27d73d243 100644
--- a/mllib-dal/src/main/scala/com/intel/oap/mllib/regression/LinearRegressionDALImpl.scala
+++ b/mllib-dal/src/main/scala/com/intel/oap/mllib/regression/LinearRegressionDALImpl.scala
@@ -106,8 +106,7 @@ class LinearRegressionDALImpl( val fitIntercept: Boolean,
     }
     lrTimer.record("Data Convertion")
 
-    val results = labeledPointsTables.mapPartitionsWithIndex {
-      case (rank: Int, tables: Iterator[(Any, Any)]) =>
+    val results = labeledPointsTables.mapPartitionsWithIndex { (rank, tables) =>
         val (feature, label) = tables.next()
         val (featureTabAddr : Long, featureRows : Long, featureColumns : Long) =
           if (useDevice == "GPU") {
diff --git a/mllib-dal/src/main/scala/com/intel/oap/mllib/regression/RandomForestRegressorDALImpl.scala b/mllib-dal/src/main/scala/com/intel/oap/mllib/regression/RandomForestRegressorDALImpl.scala
index e98a66313..03272b29b 100644
--- a/mllib-dal/src/main/scala/com/intel/oap/mllib/regression/RandomForestRegressorDALImpl.scala
+++ b/mllib-dal/src/main/scala/com/intel/oap/mllib/regression/RandomForestRegressorDALImpl.scala
@@ -75,15 +75,14 @@ class RandomForestRegressorDALImpl(val uid: String,
     }.count()
     rfrTimer.record("OneCCL Init")
 
-    val results = labeledPointsTables.mapPartitionsWithIndex {
-      (rank: Int, tables: Iterator[(String, String)]) =>
+    val results = labeledPointsTables.mapPartitionsWithIndex { (rank, tables) =>
       val (feature, label) = tables.next()
       val (featureTabAddr : Long, featureRows : Long, featureColumns : Long) = {
           val parts = feature.toString.split("_")
           (parts(0).toLong, parts(1).toLong, parts(2).toLong)
         }
       val (labelTabAddr : Long, labelRows : Long, labelColumns : Long) = {
-          val parts = feature.toString.split("_")
+          val parts = label.toString.split("_")
           (parts(0).toLong, parts(1).toLong, parts(2).toLong)
         }
       val gpuIndices = if (useDevice == "GPU") {