From 83cf258e1aabfed1fda06766a416cb358100a1ee Mon Sep 17 00:00:00 2001
From: Lucas Airam Castro de Souza <33198508+AiramL@users.noreply.github.com>
Date: Sat, 20 Jun 2020 16:28:26 -0300
Subject: [PATCH] Update dfedForest.py

---
 dfedForest.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/dfedForest.py b/dfedForest.py
index f48c7dd..2b1bae4 100644
--- 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