def METHOD_NAME(self, train_batch, batch_idx):
    """
    Training the data as batches and returns training loss on each batch
    :param train_batch: Batch data
    :param batch_idx: Batch indices
    :return: output - Training loss
    """
    x, y = train_batch
    logits = self.forward(x)
    loss = self.cross_entropy_loss(logits, y)
    return {"loss": loss}