class SimpleNet(pl.LightningModule):

    def __init__(self, label_smoothing = 0.00):
        ...
        
    def training_step(self, batch, batch_idx):
    
        # Get input and target from the batch
        x, y = batch
        
        # Run forward pass
        y_hat = self.forward(x)
        
        # Compute the loss
        loss = self.criterion(y_hat, y)
        
        # Compute metrics
        accuracy = self.accuracy(y_hat, y)
        
        # Log metrics and loss
        self.log_dict({'train_loss': loss,
                       'train_accuracy': accuracy},
                       on_step=False, on_epoch=True, prog_bar=True)
        
        return loss
ML tricks
