code: def partition(
      model, data_kwargs: dict, pp_rank: int, chunk: int, stage_num: int
  ):
      torch.manual_seed(1024)
      partition = create_partition_module(
          pp_rank, stage_num, model, data_kwargs
      )
      return partition
  
file_path: colossalai/tests/test_legacy/test_pipeline/test_middleware_1f1b.py
imports: import os
  from functools import partial
  import pytest
  import torch
  import torch.distributed.rpc
  from rpc_test_utils import DAG_MLP, MLP
  from torch._C._distributed_rpc import _is_current_rpc_agent_set
  from colossalai.fx import ColoTracer
  from colossalai.fx.passes.adding_split_node_pass import balanced_split_pass, split_with_split_nodes_pass
  from colossalai.legacy import launch
  from colossalai.legacy.pipeline.middleware.adaptor import get_fx_topology
  from colossalai.legacy.pipeline.pipeline_process_group import ppg
  from colossalai.legacy.pipeline.rpc._pipeline_schedule import OneFOneBPipelineEngine
  from colossalai.logging import disable_existing_loggers
  from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
label_desc: Use image node_id 2 for calling a global function with example usage: partition(model, data_kwargs, pp_rank, chunk, stage_num) and returns: partition