diff --git a/JSAT/src/jsat/classifiers/trees/DecisionStump.java b/JSAT/src/jsat/classifiers/trees/DecisionStump.java
index 5aa0b0f9..242ffd39 100644
--- a/JSAT/src/jsat/classifiers/trees/DecisionStump.java
+++ b/JSAT/src/jsat/classifiers/trees/DecisionStump.java
@@ -394,7 +394,7 @@ public List<List<DataPointPair<Integer>>> trainC(List<DataPointPair<Integer>> da
                 aSplit = listOfLists(catAttributes[attribute].getNumOfCategories());
                 split_scores = new ImpurityScore[aSplit.size()];
                 for(int i=0; i < split_scores.length; i++)
-                    split_scores[i] = new ImpurityScore(aSplit.size(), gainMethod);
+                    split_scores[i] = new ImpurityScore(predicting.getNumOfCategories(), gainMethod);
                 
                 List<DataPointPair<Integer>> wasMissing = new ArrayList<DataPointPair<Integer>>(0);
                 double missingSum = 0.0;