class_docstr: PipelineStageManager is a helper class to manage pipeline stages.
  
  Args:
      pg_mesh (ProcessGroupMesh): Process group mesh.
      pipeline_axis (int): The axis along which the pipeline is constructed.
      is_virtual (bool): Whether to use circle p2p communication, it will make the first and last stage communicate with each other.
  
  Attributes:
      num_stages (int): Number of stages in the pipeline.
      stage (int): The current stage.
class_docstr_tok: ["PipelineStageManager","is","a","helper","class","to","manage","pipeline","stages",".","Args",":","pg_mesh","(","ProcessGroupMesh",")",":","Process","group","mesh",".","pipeline_axis","(","int",")",":","The","axis","along","which","the","pipeline","is","constructed",".","is_virtual","(","bool",")",":","Whether","to","use","circle","p2p","communication",",","it","will","make","the","first","and","last","stage","communicate","with","each","other",".","Attributes",":","num_stages","(","int",")",":","Number","of","stages","in","the","pipeline",".","stage","(","int",")",":","The","current","stage","."]
code: def switch_model_chunk_id(self, model_chunk_id: int):
      old_model_chunk_id = self.model_chunk_id
      self.model_chunk_id = model_chunk_id
      yield
      self.model_chunk_id = old_model_chunk_id
  
file_path: colossalai/colossalai/pipeline/stage_manager.py
imports: import contextlib
  from typing import Dict, List, Optional, Tuple
  import torch.distributed
  from torch.distributed import ProcessGroup
  from colossalai.cluster import ProcessGroupMesh
label_desc: Use image node_id 11 for calling the PipelineStageManager obj's underlying member method code with example usage: obj.switch_model_chunk_id(model_chunk_id) without return types