snapshot...

modules/deeplearning/icing_fcn.py

@@ -993,7 +993,7 @@ class IcingIntensityFCN:
     def restore(self, ckpt_dir):
 
         if TRACK_MOVING_AVERAGE:
-            ckpt = tf.train.Checkpoint(step=tf.Variable(1), model=self.model, averaged_weights=self.model.trainable_variables)
+            ckpt = tf.train.Checkpoint(step=tf.Variable(1), model=self.model, model_weights=self.model.trainable_variables, averaged_weights=self.ema_trainable_variables)
         else:
             ckpt = tf.train.Checkpoint(step=tf.Variable(1), model=self.model)
 
@@ -1001,6 +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])
+
         self.reset_test_metrics()
 
         for data0, data1, label in self.test_dataset: