pl_train_epoch_end.py

import torch
import torch.nn.Functional as F
import pytorch_lightning as pl

##########################################################################################

class FlashModel(pl.LightningModule):
    """DOCSTRING"""
    def __init__(self, model):
        super().__init__()
        self.model = model
        
    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.model(x)
        loss = F.cross_entropy(y_hat, y)
        preds = ...
        return {"loss": loss, "preds": preds}
    
    def training_epoch_end(self, training_step_outputs):
        for pred in training_step_outputs:
            ## do something
            pass
        pass
    
##########################################################################################
## Under the hood pseudocode
outs = []
for batch in train_dataloader:
    ## Step-1: FORWARD
    out = training_step(val_batch)
    
    ## Step-2: BACKWARD
    loss.backward()
    
    ## Step-3: Optim step and zero-grad
    optimizer.step()
    optimizer.zero_grad()

training_epoch_end(outs)

##########################################################################################