docstr: 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
docstr_tok: ["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"]
code: def training_step(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}
  
imports: import os
  import lightning
  import torch
  from lightning.pytorch.callbacks import EarlyStopping, LearningRateMonitor, ModelCheckpoint
  from lightning.pytorch.cli import LightningCLI
  from torch.nn import functional
  from torch.utils.data import DataLoader, random_split
  from torchmetrics.functional import accuracy
  from torchvision import datasets, transforms
  import mlflow.pytorch
label_desc: Use image node_id 4 for calling the LightningMNISTClassifier obj's underlying member method code with example usage: obj.training_step(train_batch, batch_idx) and returns: dict