Fix PSNR logging.

rayanirban · Nov 4, 2021 · 60fbe22 · 60fbe22
1 parent a6bf7e4
commit 60fbe22
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/fit/modules/TRecTransformerModule.py b/fit/modules/TRecTransformerModule.py
@@ -334,8 +334,9 @@ def test_step(self, batch, batch_idx):
 
     def test_epoch_end(self, outputs):
         outputs = torch.stack(outputs)
-        self.log('Mean PSNR', torch.mean(outputs).detach().cpu().numpy(), logger=True)
-        self.log('SEM PSNR', torch.std(outputs / np.sqrt(len(outputs))).detach().cpu().numpy(),
+        print(torch.mean(outputs).detach().cpu().numpy())
+        self.log('Mean PSNR', torch.mean(outputs), logger=True)
+        self.log('SEM PSNR', torch.std(outputs / np.sqrt(len(outputs))),
                  logger=True)
 
     def get_imgs(self, x, fbp, y, amp_min, amp_max):