diff --git a/modules/deeplearning/icing_fcn.py b/modules/deeplearning/icing_fcn.py
index 382654bbf4509ac1fb953b3ab0be8dacf555470f..ca9792d7ea8d4dc8a1a9766a2c4feab306460233 100644
--- a/modules/deeplearning/icing_fcn.py
+++ b/modules/deeplearning/icing_fcn.py
@@ -992,7 +992,11 @@ class IcingIntensityFCN:
 
     def restore(self, ckpt_dir):
 
-        ckpt = tf.train.Checkpoint(step=tf.Variable(1), model=self.model)
+        if TRACK_MOVING_AVERAGE:
+            ckpt = tf.train.Checkpoint(step=tf.Variable(1), model=self.model, averaged_weights=self.model.trainable_variables)
+        else:
+            ckpt = tf.train.Checkpoint(step=tf.Variable(1), model=self.model)
+
         ckpt_manager = tf.train.CheckpointManager(ckpt, ckpt_dir, max_to_keep=3)
 
         ckpt.restore(ckpt_manager.latest_checkpoint)