code: def create_partition_module(
      pp_rank: int, stage_num: int, model, data_kwargs
  ):
      model.eval()
      tracer = ColoTracer()
      meta_args = {k: v.to("meta") for k, v in data_kwargs.items()}
      graph = tracer.trace(root=model, meta_args=meta_args)
      gm = torch.fx.GraphModule(model, graph, model.__class__.__name__)
      annotated_model = balanced_split_pass(gm, stage_num)
      top_module, split_submodules = split_with_split_nodes_pass(
          annotated_model, merge_output=True
      )
      topo = get_fx_topology(top_module)
      for submodule in split_submodules:
          if isinstance(submodule, torch.fx.GraphModule):
              setattr(submodule, "_topo", topo)
      return split_submodules[pp_rank + 1]
  
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 1 for calling a global function with example usage: create_partition_module(pp_rank, stage_num, model, data_kwargs) and returns: split_submodules