diff --git a/modules/GSOC/E2_ESRGAN/lib/train.py b/modules/GSOC/E2_ESRGAN/lib/train.py index 8a625e9253b61d1fce036ec2017ac2c54a3adc1f..d1d1fbcf75ec4b0c30b86c6987c3a69de46e640f 100644 --- a/modules/GSOC/E2_ESRGAN/lib/train.py +++ b/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)