docstr: Backward one step of the pipeline
  
  Args:
      optimizer (OptimizerWrapper): Optimizer to update the model
      input_obj (Optional[dict]): Output of the previous stage. If it is the first stage, the `input_obj` is None.
      output_obj (Union[dict, torch.Tensor]): Output of the current stage. If it is the last stage, the output is the loss (Tensor).
      output_obj_grad (dict): Gradient of the `output_obj`. If it is the last stage, the `output_obj_grad` is None.
  
  Returns:
      Optional[dict]: Gradient of the `input_obj`. If it is the first stage, the `input_obj_grad` is None.
docstr_tok: ["Backward","one","step","of","the","pipeline","Args",":","optimizer","(","OptimizerWrapper",")",":","Optimizer","to","update","the","model","input_obj","(","Optional","[","dict","]",")",":","Output","of","the","previous","stage",".","If","it","is","the","first","stage",",","the","`","input_obj","`","is","None",".","output_obj","(","Union","[","dict",",","torch.Tensor","]",")",":","Output","of","the","current","stage",".","If","it","is","the","last","stage",",","the","output","is","the","loss","(","Tensor",")",".","output_obj_grad","(","dict",")",":","Gradient","of","the","`","output_obj","`",".","If","it","is","the","last","stage",",","the","`","output_obj_grad","`","is","None",".","Returns",":","Optional","[","dict","]",":","Gradient","of","the","`","input_obj","`",".","If","it","is","the","first","stage",",","the","`","input_obj_grad","`","is","None","."]
code: def backward_step(
      self,
      optimizer: OptimizerWrapper,
      input_obj: Optional[dict],
      output_obj: Union[dict, torch.Tensor],
      output_obj_grad: Optional[dict],
  ) -> Optional[dict]:
      """Backward one step of the pipeline
  
      Args:
          optimizer (OptimizerWrapper): Optimizer to update the model
          input_obj (Optional[dict]): Output of the previous stage. If it is the first stage, the `input_obj` is None.
          output_obj (Union[dict, torch.Tensor]): Output of the current stage. If it is the last stage, the output is the loss (Tensor).
          output_obj_grad (dict): Gradient of the `output_obj`. If it is the last stage, the `output_obj_grad` is None.
  
      Returns:
          Optional[dict]: Gradient of the `input_obj`. If it is the first stage, the `input_obj_grad` is None.
      """
  
      # Retain the grad on the input_obj.
      tree_map(retain_grad, input_obj)
  
      # Backward pass.
      if output_obj_grad is None:
          optimizer.backward(output_obj)
      else:
          if "backward_tensor_keys" not in output_obj:
              for k, grad in output_obj_grad.items():
                  optimizer.backward_by_grad(output_obj[k], grad)
          else:
              for k, grad in output_obj_grad.items():
                  output_obj[k].grad = grad
              for k in output_obj["backward_tensor_keys"]:
                  tensor_to_backward = output_obj[k]
                  optimizer.backward_by_grad(
                      tensor_to_backward, tensor_to_backward.grad
                  )
  
      # Collect the grad of the input_obj.
      input_obj_grad = None
      if input_obj is not None:
          input_obj_grad = {}
          for k, v in input_obj.items():
              if isinstance(v, torch.Tensor) and v.grad is not None:
                  input_obj_grad[k] = v.grad
      return input_obj_grad
  
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 14 for calling the InterleavedSchedule obj's underlying member method code with example usage: obj.backward_step(optimizer, input_obj, output_obj, output_obj_grad) and returns: input_obj_grad