diff --git a/src/eynollah/training/utils.py b/src/eynollah/training/utils.py
index 9b4e01a..d44f10c 100644
--- a/src/eynollah/training/utils.py
+++ b/src/eynollah/training/utils.py
@@ -9,8 +9,8 @@ from scipy.ndimage.interpolation import map_coordinates
 from scipy.ndimage.filters import gaussian_filter
 from tqdm import tqdm
 import imutils
-##import tensorflow as tf
-##from tensorflow.keras.utils import to_categorical
+import tensorflow as tf
+from tensorflow.keras.utils import to_categorical
 from PIL import Image, ImageFile, ImageEnhance
 
 import torch