class_docstr: A context manager to split the model into pipeline stages.
class_docstr_tok: ["A","context","manager","to","split","the","model","into","pipeline","stages","."]
docstr: The callback function when exiting context.
docstr_tok: ["The","callback","function","when","exiting","context","."]
code: def _post_context_exec(self):
      """
      The callback function when exiting context.
      """
  
      # reset rng states
      torch.set_rng_state(self.cpu_rng_state)
      torch.cuda.set_rng_state(self.cuda_rng_state)
  
imports: import torch
  from colossalai.legacy.context import ParallelMode
  from colossalai.legacy.core import global_context
  from colossalai.legacy.nn.layer.utils import CheckpointModule
  from colossalai.tensor import ColoParameter
  from colossalai.utils.model.utils import InsertPostInitMethodToModuleSubClasses
  from .layer_spec import LayerSpec
  from .utils import build_kwargs_for_module, call_module, customized_partition, exec_funcs_with_kwargs, partition_balanced, partition_uniform
label_desc: Use image node_id 7 for calling the PipelinableContext obj's underlying member method code with example usage: obj._post_context_exec() without return types