class_docstr: Control the loops for training, evaluation, and prediction. This module is independent of
  the model definition. This class inherits from the Pytorch Lightning's LightningModule:
  https://lightning.ai/docs/pytorch/stable/common/lightning_module.html
code: def _shared_step(
      self,
      batch: Dict,
  ):
      label = batch[self.model.label_key]
      if self.mixup_fn is not None:
          self.mixup_fn.mixup_enabled = self.training & (
              self.current_epoch < self.hparams.mixup_off_epoch
          )
          batch, label = multimodel_mixup(
              batch=batch, model=self.model, mixup_fn=self.mixup_fn
          )
      output = run_model(self.model, batch)
      loss = self._compute_loss(output=output, label=label)
      return output, loss
  
imports: import logging
  from typing import Callable, Dict, List, Optional, Union
  import lightning.pytorch
  import torch
  import torch.nn.functional
  import torchmetrics
  from lightning.pytorch.strategies import DeepSpeedStrategy
  from lightning.pytorch.utilities import grad_norm
  from torch import nn
  from torch.nn.modules.loss import _Loss
  from torchmetrics.aggregation import BaseAggregator
  from ..constants import LM_TARGET, LOGITS, T_FEW, TEMPLATE_LOGITS, WEIGHT
  from ..data.mixup import MixupModule, multimodel_mixup
  from ..models.utils import run_model
  from .semantic_seg_metrics import COD, Balanced_Error_Rate
  from .utils import apply_layerwise_lr_decay, apply_single_lr, apply_two_stages_lr, get_lr_scheduler, get_optimizer
label_desc: Use image node_id 5 for calling the LitModule obj's underlying member method code with example usage: obj._shared_step(batch) and returns: output,  loss