code: def check_input(top_module, input_partition: Partition):
      partition_output = input_partition.get_output_vals()
      arg_pos = 0
      for node in top_module.graph.nodes:
          if node.op == "placeholder":
              cur_checkee = partition_output[arg_pos]
              to_partition_and_offset = cur_checkee.get()
              assert len(to_partition_and_offset) == len(
                  node.users.keys()
              )
              arg_pos += 1
  
      assert arg_pos == len(partition_output)
  
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 2 for calling a global function with example usage: check_input(top_module, input_partition) without return types