def METHOD_NAME(self, batch, batch_idx):
    out = self.forward(batch)  # PB: B, T, H
    out = out.payload if isinstance(out, PaddedBatch) else out
    labels = batch.payload
    loss_gpt = self.loss_gpt(out, labels, is_train_step=True)
    self.train_gpt_loss(loss_gpt)
    self.log(f'gpt/loss', loss_gpt, sync_dist=True)
    return loss_gpt