diff --git a/mllib/src/main/scala/org/apache/spark/mllib/tree/DecisionTree.scala b/mllib/src/main/scala/org/apache/spark/mllib/tree/DecisionTree.scala
index 6539a1194116c..9ebb1d25ffa02 100644
--- a/mllib/src/main/scala/org/apache/spark/mllib/tree/DecisionTree.scala
+++ b/mllib/src/main/scala/org/apache/spark/mllib/tree/DecisionTree.scala
@@ -941,6 +941,7 @@ object DecisionTree extends Serializable with Logging {
                 while (classIndex < numClasses) {
                   rightNodeAgg(featureIndex)(splitIndex)(classIndex)
                     = totalNodeAgg(classIndex) - leftNodeAgg(featureIndex)(splitIndex)(classIndex)
+                  classIndex += 1
                 }
                 splitIndex += 1
               }