docstr: Forward one step of the pipeline
  
  Args:
      model (Module): Model 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","(","Module",")",":","Model","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: Module,
      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 (Module): Model 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()
      # 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
      output_obj = model_forward(model, micro_batch, input_obj)
      if self.stage_manager.is_last_stage():
          loss = (
              criterion(output_obj, micro_batch) / self.num_microbatches
          )
          if accum_loss is not None:
              accum_loss.add_(loss.detach())
          if outputs is not None:
              outputs.append(tree_map_hf(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
  from torch.utils._pytree import tree_map
  from colossalai.interface import ModelWrapper, 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, tree_map_hf
  from .base import PipelineSchedule
label_desc: Use image node_id 10 for calling the OneForwardOneBackwardSchedule obj's underlying member method code with example usage: obj.forward_step(model, input_obj, criterion, accum_loss, outputs) and returns: loss, output_obj