diff --git a/mainscripts/Trainer.py b/mainscripts/Trainer.py
index fdfe386..4056fe1 100644
--- a/mainscripts/Trainer.py
+++ b/mainscripts/Trainer.py
@@ -4,6 +4,8 @@ import traceback
 import queue
 import threading
 import time
+from enum import Enum
+
 import numpy as np
 import itertools
 from pathlib import Path