diff --git a/src/AI/DecisionTrees/DecisionTree.py b/src/AI/DecisionTrees/DecisionTree.py
index 87ed395..98d2013 100644
--- a/src/AI/DecisionTrees/DecisionTree.py
+++ b/src/AI/DecisionTrees/DecisionTree.py
@@ -1,13 +1,14 @@
 from typing import List, Any
 
 from src.AI.DecisionTrees.AttributeDefinition import AttributeDefinition
+from src.AI.DecisionTrees.DecisionTreeBranch import DecisionTreeBranch
 
 
 class DecisionTree(object):
     root: AttributeDefinition
-    branches: List[Any]
+    branches: List[DecisionTreeBranch]
 
     def __init__(self, root):
         self.root = root
         self.branches = []
-        self.branchesNum = 0
\ No newline at end of file
+        self.branchesNum = 0