code: def split_model_and_get_DAG(model, data_gen):
      model.eval()
  
      # generate input sample
      kwargs = data_gen()
  
      # tracing model
      tracer = ColoTracer()
      try:
          meta_args = {k: v.to("meta") for k, v in kwargs.items()}
          graph = tracer.trace(root=model, meta_args=meta_args)
      except Exception as e:
          raise RuntimeError(
              f"Failed to trace {model.__class__.__name__}, error: {e}"
          )
      gm = GraphModule(model, graph, model.__class__.__name__)
      gm.recompile()
  
      # apply transform passes
      annotated_model = balanced_split_pass(gm, 2)
      top_module, split_submodules = split_with_split_nodes_pass(
          annotated_model
      )
  
      topo = get_fx_topology(top_module)
      for submodule in split_submodules:
          if isinstance(submodule, torch.fx.GraphModule):
              setattr(submodule, "_topo", topo)
  
      return top_module, split_submodules[0]._topo
  
file_path: ColossalAI/tests/test_fx/test_pipeline/test_topo/topo_utils.py
imports: import random
  import numpy
  import torch
  from torch.fx import GraphModule
  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.pipeline.middleware import Partition, Topo
  from colossalai.legacy.pipeline.middleware.adaptor import get_fx_topology
label_desc: Use image node_id 1 for calling a global function with example usage: split_model_and_get_DAG(model, data_gen) and returns: top_module,  split_submodules