diff --git a/modules/deeplearning/icing_cnn.py b/modules/deeplearning/icing_cnn.py
index 40af66f3f4e67b402f7ce5adc52dbd8fcfd51452..98f27613182db38f58cb7000daf1e4667b5fe180 100644
--- a/modules/deeplearning/icing_cnn.py
+++ b/modules/deeplearning/icing_cnn.py
@@ -1,5 +1,5 @@
 import tensorflow as tf
-#import tensorflow_addons as tfa
+import tensorflow_addons as tfa
 from util.setup import logdir, modeldir, cachepath
 from util.util import homedir
 
@@ -18,7 +18,7 @@ PROC_BATCH_SIZE = 2046
 PROC_BATCH_BUFFER_SIZE = 50000
 NumLabels = 1
 BATCH_SIZE = 256
-NUM_EPOCHS = 80
+NUM_EPOCHS = 50
 
 TRACK_MOVING_AVERAGE = False
 
@@ -393,7 +393,7 @@ class IcingIntensityNN:
         self.learningRateSchedule = tf.keras.optimizers.schedules.ExponentialDecay(initial_learning_rate, decay_steps, decay_rate)
 
         optimizer = tf.keras.optimizers.Adam(learning_rate=self.learningRateSchedule)
-        # optimizer = tfa.optimizers.MovingAverage(optimizer)
+        optimizer = tfa.optimizers.MovingAverage(optimizer)
 
         if TRACK_MOVING_AVERAGE:
             ema = tf.train.ExponentialMovingAverage(decay=0.999)