add maxBins to to the train method as well

maropu · May 21, 2015 · 3a09170 · 3a09170
1 parent af7f274
commit 3a09170
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/pyspark/mllib/tree.py b/python/pyspark/mllib/tree.py
@@ -435,11 +435,11 @@ class GradientBoostedTrees(object):
 
     @classmethod
     def _train(cls, data, algo, categoricalFeaturesInfo,
-               loss, numIterations, learningRate, maxDepth):
+               loss, numIterations, learningRate, maxDepth, maxBins):
         first = data.first()
         assert isinstance(first, LabeledPoint), "the data should be RDD of LabeledPoint"
         model = callMLlibFunc("trainGradientBoostedTreesModel", data, algo, categoricalFeaturesInfo,
-                              loss, numIterations, learningRate, maxDepth)
+                              loss, numIterations, learningRate, maxDepth, maxBins)
         return GradientBoostedTreesModel(model)
 
     @classmethod