diff --git a/main.py b/main.py
index b46e274..853abdb 100644
--- a/main.py
+++ b/main.py
@@ -12,12 +12,14 @@ from sklearn.tree import DecisionTreeClassifier
 
 from ClientParamsFactory import ClientParamsFactory
 from ForkliftAgent import ForkliftAgent
+from GameModel import GameModel
 from PatchAgent import PatchAgent
 from PatchType import PatchType
 from data.enum.CompanySize import CompanySize
 from data.enum.Direction import Direction
 from data.enum.Priority import Priority
 from tree.DecisionTree import DecisionTree
+from util.PathDefinitions import GridWithWeights
 
 colors = [
     'blue', 'cyan', 'orange', 'yellow', 'magenta', 'purple', '#103d3e', '#9fc86c',
diff --git a/pathfinding/PathFinderState.py b/pathfinding/PathFinderState.py
index 4195016..832ac85 100644
--- a/pathfinding/PathFinderState.py
+++ b/pathfinding/PathFinderState.py
@@ -1,6 +1,6 @@
 from typing import List
 
-from data.Direction import Direction
+from data.enum.Direction import Direction
 from decision.ActionType import ActionType
 from util.PathDefinitions import GridLocation