diff --git a/osrt/model.py b/osrt/model.py
index 9e1e28a0ead36bf766c9874003ff0a2dadbae95e..ebd7a171fda91fa085f68bff4f0d44311b33aad8 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': loss_value.item()}
+        return {'loss': loss_value}
     
     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': loss_value.item(), 'val_psnr': psnr.item()}
+        return {'loss': loss_value, 'val_psnr': psnr.item()}