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
class_docstr_tok: ["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.
      """
      if self.signature == QUERY:
          embeddings = run_model(self.query_model, batch)[
              self.query_model.prefix
          ][FEATURES]
          return {FEATURES: embeddings}
      elif self.signature == RESPONSE:
          embeddings = run_model(self.response_model, batch)[
              self.response_model.prefix
          ][FEATURES]
          return {FEATURES: embeddings}
      else:
          query_embeddings = run_model(self.query_model, batch)[
              self.query_model.prefix
          ][FEATURES]
          response_embeddings = run_model(self.response_model, batch)[
              self.response_model.prefix
          ][FEATURES]
  
      match_prob = compute_matching_probability(
          embeddings1=query_embeddings,
          embeddings2=response_embeddings,
      )
      if self.match_label == 0:
          probability = torch.stack([match_prob, 1 - match_prob]).t()
      else:
          probability = torch.stack([1 - match_prob, match_prob]).t()
  
      return {PROBABILITY: probability}
  
imports: import logging
  from typing import Callable, Dict, List, Optional, Union
  import lightning.pytorch
  import torch
  import torchmetrics
  from lightning.pytorch.utilities import grad_norm
  from omegaconf import DictConfig
  from torch import nn
  from torch.nn.modules.loss import _Loss
  from torchmetrics.aggregation import BaseAggregator
  from ..constants import FEATURES, LOGIT_SCALE, PROBABILITY, QUERY, RESPONSE
  from ..models.utils import run_model
  from ..utils.matcher import compute_matching_probability
  from .losses import MultiNegativesSoftmaxLoss
  from .utils import CustomHitRate, apply_layerwise_lr_decay, apply_single_lr, apply_two_stages_lr, generate_metric_learning_labels, get_lr_scheduler, get_optimizer
label_desc: Use image node_id 8 for calling the MatcherLitModule obj's underlying member method code with example usage: obj.predict_step(batch, batch_idx, dataloader_idx) and returns: dict, dict, dict