diff --git a/generator/ctrl/training_utils/training.py b/generator/ctrl/training_utils/training.py
index a59965c..0cb78a2 100644
--- a/generator/ctrl/training_utils/training.py
+++ b/generator/ctrl/training_utils/training.py
@@ -18,6 +18,7 @@ from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import embedding_ops
 import fastBPE
 import platform
+from tensorflow.keras.callbacks import TensorBoard, EarlyStopping
 
 use_py3 = platform.python_version()[0] == '3'