class_docstr: The lightning module for a merged "multi-model".
  
  The output of the multi-model is expected to be a tuple of tensors.
  The tensors will be each passed to a criterion and a metric.
  The loss will be added up for back propagation, and the metrics will be logged.
  
  The reported metric will be a list of metrics, one for each model.
  
  Parameters
  ----------
  criterion
      Loss function.
  metric
      Metric function.
  n_models
      Number of models in the multi-model.
class_docstr_tok: ["The","lightning","module","for","a","merged","``","multi-model","''",".","The","output","of","the","multi-model","is","expected","to","be","a","tuple","of","tensors",".","The","tensors","will","be","each","passed","to","a","criterion","and","a","metric",".","The","loss","will","be","added","up","for","back","propagation",",","and","the","metrics","will","be","logged",".","The","reported","metric","will","be","a","list","of","metrics",",","one","for","each","model",".","Parameters","--","--","--","--","--","criterion","Loss","function",".","metric","Metric","function",".","n_models","Number","of","models","in","the","multi-model","."]
code: def training_step(self, batch, batch_idx):
      x, y = batch
      multi_y_hat = self(x)
      if isinstance(multi_y_hat, tuple):
          assert len(multi_y_hat) == self.n_models
      else:
          assert self.n_models == 1
          multi_y_hat = [multi_y_hat]
      multi_loss = []
      for idx, y_hat in enumerate(multi_y_hat):
          loss = self.criterion(y_hat.to("cpu"), y.to("cpu"))
          self.log(f"train_loss_{idx}", loss, prog_bar=True)
          self.log(
              f"train_metric_{idx}",
              self.metric(y_hat.to("cpu"), y.to("cpu")),
              prog_bar=True,
          )
          multi_loss.append(loss)
      return sum(multi_loss)
  
imports: from __future__ import annotations
  import pytorch_lightning
  import torch
  from pytorch_lightning.strategies import SingleDeviceStrategy
  from torch import nn
  from torchmetrics import Metric
  import nni
  from nni.nas.evaluator.pytorch.lightning import LightningModule
label_desc: Use image node_id 5 for calling the MultiModelLightningModule obj's underlying member method code with example usage: obj.training_step(batch, batch_idx) and returns: sum