docstr: Args:
      model_chunk (ModuleList or Module): Model Chunk to be trained. Original interleaved uses a module list whereas shardformer uses entire model + layer specification
      data_iter (Iterable): Data iterator.
      criterion (Callable[[Any, Any], Tensor]): Criterion to be used. It should take two arguments: model outputs and inputs, and returns loss tensor.
      optimizer (OptimizerWrapper, optional): Optimizer to be used. Can be None when only forward is executed. Defaults to None.
      return_loss (bool, optional): Whether to return loss. Defaults to False. Whether to return loss.
      return_outputs (bool, optional): Whether to return model outputs. Defaults to False. Whether to return model outputs.
  
  Returns:
      dict: A dict with keys: 'loss' and 'outputs'.
docstr_tok: ["Args",":","model_chunk","(","ModuleList","or","Module",")",":","Model","Chunk","to","be","trained",".","Original","interleaved","uses","a","module","list","whereas","shardformer","uses","entire","model","+","layer","specification","data_iter","(","Iterable",")",":","Data","iterator",".","criterion","(","Callable","[","[","Any",",","Any","]",",","Tensor","]",")",":","Criterion","to","be","used",".","It","should","take","two","arguments",":","model","outputs","and","inputs",",","and","returns","loss","tensor",".","optimizer","(","OptimizerWrapper",",","optional",")",":","Optimizer","to","be","used",".","Can","be","None","when","only","forward","is","executed",".","Defaults","to","None",".","return_loss","(","bool",",","optional",")",":","Whether","to","return","loss",".","Defaults","to","False",".","Whether","to","return","loss",".","return_outputs","(","bool",",","optional",")",":","Whether","to","return","model","outputs",".","Defaults","to","False",".","Whether","to","return","model","outputs",".","Returns",":","dict",":","A","dict","with","keys",":","'loss","'","and","'outputs","'","."]
code: def forward_backward_step(
      self,
      model_chunk: Union[ModuleList, Module],
      data_iter: Iterable,
      criterion: Callable[..., Any],
      optimizer: Optional[OptimizerWrapper] = None,
      return_loss: bool = False,
      return_outputs: bool = False,
  ) -> dict:
      """
      Args:
          model_chunk (ModuleList or Module): Model Chunk to be trained. Original interleaved uses a module list whereas shardformer uses entire model + layer specification
          data_iter (Iterable): Data iterator.
          criterion (Callable[[Any, Any], Tensor]): Criterion to be used. It should take two arguments: model outputs and inputs, and returns loss tensor.
          optimizer (OptimizerWrapper, optional): Optimizer to be used. Can be None when only forward is executed. Defaults to None.
          return_loss (bool, optional): Whether to return loss. Defaults to False. Whether to return loss.
          return_outputs (bool, optional): Whether to return model outputs. Defaults to False. Whether to return model outputs.
  
      Returns:
          dict: A dict with keys: 'loss' and 'outputs'.
      """
      self.forward_only = not torch.is_grad_enabled()
      if optimizer is None:
          assert (
              self.forward_only
          ), "Optimizer should be passed when doing backward."
  
      if self.forward_only:
          result = self.run_forward_only(
              model_chunk,
              data_iter,
              criterion,
              return_loss,
              return_outputs,
          )
      else:
          result = self.run_forward_backward(
              model_chunk,
              data_iter,
              criterion,
              optimizer,
              return_loss,
              return_outputs,
          )
  
      return result
  
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 17 for calling the InterleavedSchedule obj's underlying member method code with example usage: obj.forward_backward_step(model_chunk, data_iter, criterion, optimizer, return_loss, return_outputs) and returns: result