diff --git a/mllib-dal/src/main/scala/com/intel/oap/mllib/OneDAL.scala b/mllib-dal/src/main/scala/com/intel/oap/mllib/OneDAL.scala
index c89b63463..cde788ef8 100644
--- a/mllib-dal/src/main/scala/com/intel/oap/mllib/OneDAL.scala
+++ b/mllib-dal/src/main/scala/com/intel/oap/mllib/OneDAL.scala
@@ -501,7 +501,7 @@ object OneDAL {
     }
 
     // Get dimensions for each partition
-    val partitionDims = Utils.getPartitionDims(labeledPoints.select(featuresCol).rdd.map { row =>
+    val partitionDims = Utils.getPartitionDims(dataForConversion.select(featuresCol).rdd.map{ row =>
       val vector = row.getAs[Vector](0)
       vector
     })
@@ -510,7 +510,7 @@ object OneDAL {
     // "No partitions or no locations for partitions found".
     // TODO: ML-312: Improve ExecutorInProcessCoalescePartitioner
     val nonEmptyPartitions = dataForConversion.select(labelCol, featuresCol).toDF().rdd.mapPartitionsWithIndex {
-      (index: Int, it: Iterator[Row]) => Iterator(Tuple3( partitionDims(index)._1,index, it))
+      (index: Int, it: Iterator[Row]) => Iterator(Tuple3(partitionDims(index)._1, index, it))
     }.filter {
       _._1 > 0
     }.flatMap {