code: def check_submod(top_module, part_id, mid_partition: Partition):
      partition_input = mid_partition.get_input_vals()
      partition_output = mid_partition.get_output_vals()
  
      cnt = 1
      cur_node = None
      for node in top_module.graph.nodes:
          if node.name.startswith("submod"):
              cnt += 1
          if cnt == part_id:
              cur_node = node
              break
  
      assert len(partition_input) == len(cur_node.args)
      assert len(partition_output) == len(cur_node.users)
  
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 3 for calling a global function with example usage: check_submod(top_module, part_id, mid_partition) without return types