docstr: Copy the gradient tensor from the next stage in pipeline as the input gradient of this stage.
     For interleaved 1F1B.
  
  Args:
      model_chunk_id (int): The current model chunk idx.
      next_rank (int, optional): The rank of the source of the tensor.
  
  Returns:
      Any: The input gradient tensor or gradient tensor list.
docstr_tok: ["Copy","the","gradient","tensor","from","the","next","stage","in","pipeline","as","the","input","gradient","of","this","stage",".","For","interleaved","1F1B",".","Args",":","model_chunk_id","(","int",")",":","The","current","model","chunk","idx",".","next_rank","(","int",",","optional",")",":","The","rank","of","the","source","of","the","tensor",".","Returns",":","Any",":","The","input","gradient","tensor","or","gradient","tensor","list","."]
code: def recv_backward(
      self, model_chunk_id: int, next_rank: int = None
  ) -> Any:
      """Copy the gradient tensor from the next stage in pipeline as the input gradient of this stage.
         For interleaved 1F1B.
  
      Args:
          model_chunk_id (int): The current model chunk idx.
          next_rank (int, optional): The rank of the source of the tensor.
  
      Returns:
          Any: The input gradient tensor or gradient tensor list.
      """
      with self.stage_manager.switch_model_chunk_id(model_chunk_id):
          if not self.stage_manager.is_last_stage():
              output_tensor_grad = self.comm.recv_backward(
                  next_rank, metadata_recv=self.grad_metadata_recv
              )
              if (
                  self.enable_metadata_cache
                  and self.grad_metadata_recv is None
              ):
                  self.grad_metadata_recv = create_send_metadata(
                      output_tensor_grad
                  )
  
              return output_tensor_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 6 for calling the InterleavedSchedule obj's underlying member method code with example usage: obj.recv_backward(model_chunk_id, next_rank) and returns: output_tensor_grad