From db4d7175ca1bbc4d4842b3e7d9855f0d4d8234ac Mon Sep 17 00:00:00 2001 From: alexcbb <alexchapin@hotmail.fr> Date: Tue, 25 Jul 2023 10:30:19 +0200 Subject: [PATCH] Adapt return of steps --- osrt/model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/osrt/model.py b/osrt/model.py index d989406..bc816da 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 loss_value.item() + return {'train_mse': 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 loss_value.item(), psnr.item() + return {'val_mse': loss_value.item(), 'val_psnr': psnr.item()} -- GitLab