code: def avgnode_split_pass(gm: torch.fx.GraphModule, pp_size: int):
      """
      In avgnode_split_pass, simply split graph by node number.
      """
      mod_graph = gm.graph
      avg_num_node = len(mod_graph.nodes) // pp_size
      accumulate_num_node = 0
      for node in mod_graph.nodes:
          if pp_size <= 1:
              break
          accumulate_num_node += 1
          if accumulate_num_node >= avg_num_node:
              accumulate_num_node = 0
              pp_size -= 1
              if node.next.op == "output":
                  with mod_graph.inserting_before(node):
                      split_node = mod_graph.create_node(
                          "call_function", pipe_split
                      )
              else:
                  with mod_graph.inserting_after(node):
                      split_node = mod_graph.create_node(
                          "call_function", pipe_split
                      )
      gm.recompile()
      return gm
  
file_path: ColossalAI/colossalai/fx/passes/adding_split_node_pass.py
imports: import numpy
  import torch
  import tqdm
  from colossalai.fx.passes.split_module import split_module
label_desc: Use image node_id 10 for calling a global function with example usage: avgnode_split_pass(gm, pp_size) and returns: gm