snapshot...

modules/GSOC/E2_ESRGAN/lib/train.py

@@ -71,6 +71,7 @@ class Trainer(object):
     total_steps = phase_args["num_steps"]
     metric = tf.keras.metrics.Mean()
     psnr_metric = tf.keras.metrics.Mean()
+    mse_metric = tf.keras.metrics.Mean()
 
     # Generator Optimizer
     G_optimizer = tf.optimizers.Adam(
@@ -117,6 +118,7 @@ class Trainer(object):
     for epoch in range(NUM_WU_EPOCHS):
         print('start epoch #: ', epoch)
         metric.reset_states()
+        psnr_metric.reset_states()
         for image_lr, image_hr in self.dataset:
           num_steps = train_step(image_lr, image_hr)