def METHOD_NAME(self, batch, batch_idx):
    x, y = batch
    y_hat = self(x)
    loss = self.criterion(y_hat, y)
    self.log('train_loss', loss, prog_bar=True)
    self.log('train_accuracy', self.accuracy(y_hat, y), prog_bar=True)
    return loss