code: def build_pipeline(model):
      from colossalai.legacy.pipeline.utils import partition_uniform
  
      pipeline_size = gpc.get_world_size(ParallelMode.PIPELINE)
      pipeline_rank = gpc.get_local_rank(ParallelMode.PIPELINE)
      depth = len(model)
      start, end = partition_uniform(depth, pipeline_size, 1)[
          pipeline_rank
      ][0]
      layers = []
      for i in range(depth):
          if start <= i < end:
              layers.append(model[i])
          else:
              layers.append(nn.Identity())
      return nn.Sequential(*tuple(layers))
  
file_path: colossalai/tests/test_legacy/test_utils/test_checkpoint/test_checkpoint_2d.py
imports: import pprint
  import pytest
  import torch
  import torch.nn
  import colossalai.legacy.nn
  from colossalai.legacy.context.parallel_mode import ParallelMode
  from colossalai.legacy.core import global_context
  from colossalai.legacy.initialize import launch
  from colossalai.legacy.utils import is_using_pp
  from colossalai.legacy.utils.checkpointing import gather_pipeline_parallel_state_dict, load_checkpoint, save_checkpoint
  from colossalai.logging import disable_existing_loggers
  from colossalai.testing import rerun_if_address_is_in_use, skip_if_not_enough_gpus, spawn
label_desc: Use image node_id 1 for calling a global function with example usage: build_pipeline(model) and returns: nn