code: def check_layer(rank, world_size, port):
      disable_existing_loggers()
      launch(
          config=CONFIG,
          rank=rank,
          world_size=world_size,
          host="localhost",
          port=port,
          backend="nccl",
      )
      input_tensor = torch.rand(2, 16).cuda()
      model = MLP(16).cuda()
      symbolic_traced = symbolic_trace(model)
      output = model(input_tensor)
      splitted_gm = column_shard_linear_pass(symbolic_traced)
      new_output = splitted_gm(input_tensor)
  
      assert output.equal(new_output)
  
      gpc.destroy()
      torch.cuda.empty_cache()
  
file_path: colossalai/tests/test_fx/test_parallel_1d.py
imports: import pytest
  import torch
  from torch.fx import symbolic_trace
  from colossalai.fx.passes import column_shard_linear_pass
  from colossalai.initialize import launch
  from colossalai.legacy.core import global_context
  from colossalai.logging import disable_existing_loggers
  from colossalai.testing import clear_cache_before_run, rerun_if_address_is_in_use, spawn
label_desc: Use image node_id 1 for calling a global function with example usage: check_layer(rank, world_size, port) without return types