Update dfedForest.py

GTA-UFRJ · Jun 20, 2020 · 83cf258 · 83cf258
1 parent 1808dc3
commit 83cf258
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/dfedForest.py b/dfedForest.py
@@ -61,8 +61,8 @@ def loadForest(self):
                 self.forestList.append(loads(readFile.read()))            
 
     # Generate a new tree based on a dataset
-    def createTree(self):
-        self.newTree = tree.DecisionTreeRegressor()
+    def createTree(self,depth):
+        self.newTree = tree.DecisionTreeClassifier(max_depth=depth)
         self.newTree = self.newTree.fit(self.data, self.label)
 
     # Show all tree on the list