docstr: Forward one step of the pipeline
  Args:
      model (ModuleList or Module): Model Chunk to be run
      input_obj (Optional[dict]): The output from the previous stage. If it is the first stage, the `input_obj` is None.
      criterion (Callable): Criterion to calculate loss.
      accum_loss (Optional[torch.Tensor], optional): Accumulated loss. Defaults to None.
      outputs (Optional[List[Any]], optional): List to store the output of the last stage (final output). Defaults to None.
  
  Returns:
      Union[torch.Tensor, dict]: The intermediate output (dict) of the current stage. If it is the last stage, the output is the loss (Tensor).
docstr_tok: ["Forward","one","step","of","the","pipeline","Args",":","model","(","ModuleList","or","Module",")",":","Model","Chunk","to","be","run","input_obj","(","Optional","[","dict","]",")",":","The","output","from","the","previous","stage",".","If","it","is","the","first","stage",",","the","`","input_obj","`","is","None",".","criterion","(","Callable",")",":","Criterion","to","calculate","loss",".","accum_loss","(","Optional","[","torch.Tensor","]",",","optional",")",":","Accumulated","loss",".","Defaults","to","None",".","outputs","(","Optional","[","List","[","Any","]","]",",","optional",")",":","List","to","store","the","output","of","the","last","stage","(","final","output",")",".","Defaults","to","None",".","Returns",":","Union","[","torch.Tensor",",","dict","]",":","The","intermediate","output","(","dict",")","of","the","current","stage",".","If","it","is","the","last","stage",",","the","output","is","the","loss","(","Tensor",")","."]
code: def forward_step(
      self,
      model_chunk: Union[ModuleList, Module],
      model_chunk_id: int,
      input_obj: Optional[dict],
      criterion: Callable,
      accum_loss: Optional[torch.Tensor] = None,
      outputs: Optional[List[Any]] = None,
  ) -> Union[torch.Tensor, dict]:
      """Forward one step of the pipeline
      Args:
          model (ModuleList or Module): Model Chunk to be run
          input_obj (Optional[dict]): The output from the previous stage. If it is the first stage, the `input_obj` is None.
          criterion (Callable): Criterion to calculate loss.
          accum_loss (Optional[torch.Tensor], optional): Accumulated loss. Defaults to None.
          outputs (Optional[List[Any]], optional): List to store the output of the last stage (final output). Defaults to None.
  
      Returns:
          Union[torch.Tensor, dict]: The intermediate output (dict) of the current stage. If it is the last stage, the output is the loss (Tensor).
      """
      micro_batch = self.load_micro_batch(model_chunk_id=model_chunk_id)
  
      # for the first stage, input_obj is None
      # for the non-first stage, input_obj is the output of the previous stage and it's must be a dict
  
      with self.stage_manager.switch_model_chunk_id(model_chunk_id):
          if isinstance(model_chunk, ModuleList):
              output_obj = model_forward(
                  model_chunk[model_chunk_id], micro_batch, input_obj
              )
          else:
              # NOTE: in shardformer, each device still has the entire model, so we need to use relevant stage layers
              internal_inputs = {} if input_obj is None else input_obj
              internal_inputs[
                  "stage_index"
              ] = self.stage_manager.stage_indices[model_chunk_id]
              output_obj = model_forward(
                  model_chunk, micro_batch, internal_inputs
              )
  
          if self.stage_manager.is_last_stage():
              loss = (
                  criterion(output_obj, micro_batch)
                  / self.num_microbatch
              )
              if accum_loss is not None:
                  accum_loss.add_(loss.detach())
              if outputs is not None:
                  outputs.append(tree_map(detach, output_obj))
              return loss
          else:
              return output_obj
  
imports: from functools import partial
  from typing import Any, Callable, Dict, Iterable, List, Optional, Union
  import torch
  import torch.cuda
  from torch.nn import Module, ModuleList
  from torch.utils._pytree import tree_map
  from colossalai.interface import OptimizerWrapper
  from colossalai.pipeline.p2p import PipelineP2PCommunication, create_send_metadata
  from colossalai.pipeline.stage_manager import PipelineStageManager
  from colossalai.utils.device import get_current_device
  from ._utils import detach, get_batch_size, get_micro_batch, merge_batch, model_forward, retain_grad, to_device
  from .base import PipelineSchedule
label_desc: Use image node_id 13 for calling the InterleavedSchedule obj's underlying member method code with example usage: obj.forward_step(model_chunk, model_chunk_id, input_obj, criterion, accum_loss, outputs) and returns: loss, output_obj