diff --git a/code/python/main.py b/code/python/main.py
index 27e43dd..e7a4dac 100644
--- a/code/python/main.py
+++ b/code/python/main.py
@@ -13,6 +13,7 @@ sys.path.insert(0, os.getcwd());
 
 from src.thirdparty.maths import *;
 
+from models.generated.config import *;
 from models.generated.commands import *;
 from src.core.log import *;
 from src.setup.config import *;
@@ -36,7 +37,7 @@ def enter():
         if isinstance(command, CommandTsp):
             tsp_algorithm(
                 dist = np.asarray(command.dist, dtype=float),
-                optimise = min if command.optimise == EnumTspOptimise.min else max,
+                optimise = min if command.optimise == EnumTSPOptimise.min else max,
                 verbose = OPTIONS.tsp.verbose,
             );
         elif isinstance(command, CommandHirschberg):