diff --git a/modules/deeplearning/icing_fcn.py b/modules/deeplearning/icing_fcn.py
index 1929d96e3d2f60f9b96b1a336e7e9a367004648a..0905a2e89ac82f2141b0b97f5c26cba45bf6e57a 100644
--- a/modules/deeplearning/icing_fcn.py
+++ b/modules/deeplearning/icing_fcn.py
@@ -1001,8 +1001,9 @@ class IcingIntensityFCN:
 
         ckpt.restore(ckpt_manager.latest_checkpoint)
 
-        for idx, var in enumerate(self.model.trainable_variables):
-            var.assign(self.ema_trainable_variables[idx])
+        if TRACK_MOVING_AVERAGE:
+            for idx, var in enumerate(self.model.trainable_variables):
+                var.assign(self.ema_trainable_variables[idx])
 
         self.reset_test_metrics()