From 0089a34e6282b04b205a7514403458063fb276d5 Mon Sep 17 00:00:00 2001
From: tomrink <rink@ssec.wisc.edu>
Date: Tue, 25 Oct 2022 14:59:50 -0500
Subject: [PATCH] snapshot...

---
 modules/deeplearning/icing_fcn.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/modules/deeplearning/icing_fcn.py b/modules/deeplearning/icing_fcn.py
index ca9792d7..1b6d7e06 100644
--- a/modules/deeplearning/icing_fcn.py
+++ b/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:
-- 
GitLab