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
docstr: Per prediction step. This function is registered by LightningModule.
  Refer to https://lightning.ai/docs/pytorch/stable/common/lightning_module.html#prediction-loop
  
  Parameters
  ----------
  batch
      A dictionary containing the mini-batch data.
      The mini-batch data are passed to each individual model,
      which indexes its required input data by keys with its model prefix.
      Ground-truth labels are not needed for prediction.
  batch_idx
      Index of mini-batch.
  dataloader_idx
      Index of dataloader.
  Returns
  -------
  A dictionary with the mini-batch's logits and features.
docstr_tok: ["Per","prediction","step",".","This","function","is","registered","by","LightningModule",".","Refer","to","https",":","\/\/lightning.ai\/docs\/pytorch\/stable\/common\/lightning_module.html","#","prediction-loop","Parameters","--","--","--","--","--","batch","A","dictionary","containing","the","mini-batch","data",".","The","mini-batch","data","are","passed","to","each","individual","model",",","which","indexes","its","required","input","data","by","keys","with","its","model","prefix",".","Ground-truth","labels","are","not","needed","for","prediction",".","batch_idx","Index","of","mini-batch",".","dataloader_idx","Index","of","dataloader",".","Returns","--","--","--","-","A","dictionary","with","the","mini-batch","'s","logits","and","features","."]
code: def predict_step(self, batch, batch_idx, dataloader_idx=0):
      """
      Per prediction step. This function is registered by LightningModule.
      Refer to https://lightning.ai/docs/pytorch/stable/common/lightning_module.html#prediction-loop
  
      Parameters
      ----------
      batch
          A dictionary containing the mini-batch data.
          The mini-batch data are passed to each individual model,
          which indexes its required input data by keys with its model prefix.
          Ground-truth labels are not needed for prediction.
      batch_idx
          Index of mini-batch.
      dataloader_idx
          Index of dataloader.
      Returns
      -------
      A dictionary with the mini-batch's logits and features.
      """
      output = run_model(self.model, batch)
      if self.model_postprocess_fn:
          output = self.model_postprocess_fn(output)
  
      return output[self.model.prefix]
  
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 9 for calling the LitModule obj's underlying member method code with example usage: obj.predict_step(batch, batch_idx, dataloader_idx) and returns: output