code: def training_step(self, batch, batch_idx):
      x, y = batch
      logits = self.forward(x)
      loss = self.cross_entropy_loss(logits, y)
      self.train_acc(torch.argmax(logits, dim=1), y)
      self.log(
          "train_acc",
          self.train_acc.compute(),
          on_step=False,
          on_epoch=True,
      )
      self.log("loss", loss)
      return {"loss": loss}
  
file_path: mlflow/examples/pytorch/AxHyperOptimizationPTL/iris.py
imports: import pytorch_lightning
  import torch
  import torch.nn.functional
  from torch import nn
label_desc: Use image node_id 4 for calling the IrisClassification obj's underlying member method code with example usage: obj.training_step(batch, batch_idx) and returns: dict