Skip to content
Snippets Groups Projects
Commit db4d7175 authored by Alexandre Chapin's avatar Alexandre Chapin :race_car:
Browse files

Adapt return of steps

parent 9103c987
No related branches found
No related tags found
No related merge requests found
......@@ -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()}
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment