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","."]
code: def __init__(self, policy: str = "balanced"):
      super().__init__()
      self._layer_spec_dict = {}
      self._root_children = None
      self._model = None
      self._layer_spec_list = []
      self._func_dict = {}
      self._policy = policy
  
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 1 to create a new PipelinableContext object from inherited base classes: InsertPostInitMethodToModuleSubClasses with example: obj = PipelinableContext(policy)