Update main_train_psnr.py

haoyuc · Jul 22, 2023 · b4cb8ed · b4cb8ed
1 parent d7d50ce
commit b4cb8ed
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/main_train_psnr.py b/main_train_psnr.py
@@ -326,6 +326,8 @@ def main(json_path='options/train_msrresnet_psnr.json'):
 
                 writer.add_scalar('PSNRY', avg_psnrY, global_step=current_step)
                 writer.add_scalar('SSIMY', avg_ssimY, global_step=current_step)
+                writer.add_scalar('PSNR',  avg_psnr,  global_step=current_step)
+                writer.add_scalar('SSIM',  avg_ssim,  global_step=current_step)
                 writer.add_scalar('LPIPS', avg_lpips, global_step=current_step)