diff --git a/modules/deeplearning/icing_fcn.py b/modules/deeplearning/icing_fcn.py
index f70d69d73e4e98e491c1a66f2685839c60c9f611..015dd0474538f192a20e9bb6d6495a0c0772702c 100644
--- a/modules/deeplearning/icing_fcn.py
+++ b/modules/deeplearning/icing_fcn.py
@@ -316,26 +316,26 @@ class IcingIntensityFCN:
             label = np.where(np.invert(np.logical_or(label == 0, label == 1)), 2, label)
             label = label.reshape((label.shape[0], 1))
 
-        # if is_training and DO_AUGMENT:
-        #     data_ud = np.flip(data, axis=1)
-        #     data_alt_ud = np.copy(data_alt)
-        #     label_ud = np.copy(label)
-        #
-        #     data_lr = np.flip(data, axis=2)
-        #     data_alt_lr = np.copy(data_alt)
-        #     label_lr = np.copy(label)
-        #
-        #     data_r1 = np.rot90(data, k=1, axes=(1, 2))
-        #     data_alt_r1 = np.copy(data_alt)
-        #     label_r1 = np.copy(label)
-        #
-        #     data_r2 = np.rot90(data, k=1, axes=(1, 2))
-        #     data_alt_r2 = np.copy(data_alt)
-        #     label_r2 = np.copy(label)
-        #
-        #     data = np.concatenate([data, data_ud, data_lr, data_r1, data_r2])
-        #     data_alt = np.concatenate([data_alt, data_alt_ud, data_alt_lr, data_alt_r1, data_alt_r2])
-        #     label = np.concatenate([label, label_ud, label_lr, label_r1, label_r2])
+        if is_training and DO_AUGMENT:
+            data_ud = np.flip(data, axis=1)
+            data_alt_ud = np.copy(data_alt)
+            label_ud = np.copy(label)
+
+            data_lr = np.flip(data, axis=2)
+            data_alt_lr = np.copy(data_alt)
+            label_lr = np.copy(label)
+
+            data_r1 = np.rot90(data, k=1, axes=(1, 2))
+            data_alt_r1 = np.copy(data_alt)
+            label_r1 = np.copy(label)
+
+            data_r2 = np.rot90(data, k=1, axes=(1, 2))
+            data_alt_r2 = np.copy(data_alt)
+            label_r2 = np.copy(label)
+
+            data = np.concatenate([data, data_ud, data_lr, data_r1, data_r2])
+            data_alt = np.concatenate([data_alt, data_alt_ud, data_alt_lr, data_alt_r1, data_alt_r2])
+            label = np.concatenate([label, label_ud, label_lr, label_r1, label_r2])
 
         return data, data_alt, label
 
@@ -477,8 +477,8 @@ class IcingIntensityFCN:
         dataset = dataset.map(self.data_function, num_parallel_calls=8)
         dataset = dataset.cache()
         dataset = dataset.shuffle(PROC_BATCH_BUFFER_SIZE, reshuffle_each_iteration=True)
-        if DO_AUGMENT:
-            dataset = dataset.map(augment_image(), num_parallel_calls=8)
+        # if DO_AUGMENT:
+        #     dataset = dataset.map(augment_image(), num_parallel_calls=8)
         dataset = dataset.prefetch(buffer_size=1)
         self.train_dataset = dataset