From cf2259ac8aa087d1f00f7c1f6ddfc912f4b20b9a Mon Sep 17 00:00:00 2001
From: alexcbb <alexchapin@hotmail.fr>
Date: Tue, 25 Jul 2023 10:33:52 +0200
Subject: [PATCH] Return right format output
---
osrt/model.py | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/osrt/model.py b/osrt/model.py
index bc816da..9e1e28a 100644
--- a/osrt/model.py
+++ b/osrt/model.py
@@ -172,7 +172,7 @@ class LitSlotAttentionAutoEncoder(pl.LightningModule):
self.optimizer.step()
self.log('train_mse', loss_value, on_epoch=True)
- return {'train_mse': loss_value.item()}
+ return {'loss': loss_value.item()}
def validation_step(self, batch, batch_idx):
"""Perform a single eval step."""
@@ -189,5 +189,5 @@ class LitSlotAttentionAutoEncoder(pl.LightningModule):
self.log('val_mse', loss_value)
self.log('val_psnr', psnr)
- return {'val_mse': loss_value.item(), 'val_psnr': psnr.item()}
+ return {'loss': loss_value.item(), 'val_psnr': psnr.item()}
--
GitLab