pytorch-lightning/tests/base/eval_model_train_steps.py

from abc import ABC
from collections import OrderedDict


class TrainingStepVariations(ABC):
    """
    Houses all variations of training steps
    """
    def training_step(self, batch, batch_idx, optimizer_idx=None):
        """Lightning calls this inside the training loop"""
        # forward pass
        x, y = batch
        x = x.view(x.size(0), -1)

        y_hat = self(x)

        # calculate loss
        loss_val = self.loss(y, y_hat)

        # alternate possible outputs to test
        if self.trainer.batch_idx % 1 == 0:
            output = OrderedDict({
                'loss': loss_val,
                'progress_bar': {'some_val': loss_val * loss_val},
                'log': {'train_some_val': loss_val * loss_val},
            })
            return output

        if self.trainer.batch_idx % 2 == 0:
            return loss_val