MNISTClassifier_06.py

    def training_step(self, batch, batch_idx):
        '''Called for every training step, uses NLL Loss to compute training loss, then logs and sends back 
        logs parameter to Trainer to perform backpropagation
        '''
        
        # Get input and output from batch
        x, labels = batch
        
        # Compute prediction through the network
        prediction = self.forward(x)
        
        loss = F.nll_loss(prediction, labels)
        
        # Logs training loss
        logs={'train_loss':loss}
        
        output = {
            # This is required in training to be used by backpropagation
            'loss':loss,
            # This is optional for logging pourposes
            'log':logs
        }
        
        return output