Commit local changes (there are a lot)

lightning_models.py

@@ -161,7 +161,7 @@ class Seq2SeqLightning(LightningModule):
         output = self(input_ids, attention_mask, decoder_input_ids, decoder_attention_mask)
         loss = self.criterion(output.view(-1, output.size(-1)), labels.view(-1), ignore_index=0)
         self.log('train_loss', loss)
-        return loss
+        return loss.item()
 
     def validation_step(self, batch, batch_idx):
         input_ids = batch['input_ids']