code: def check_mem():
      pg = ProcessGroup(tp_degree=dist.get_world_size())
      size = dist.get_world_size()
      assert torch.cuda.memory_allocated() == 0
      x = torch.rand(32, 32).cuda()
      orig_mem = x.numel() * x.element_size()
      assert torch.cuda.memory_allocated() == orig_mem
      old_dist_spec = ReplicaSpec()
      row_spec = ShardSpec([0], [size])
      x.data = DistSpecManager._shard_as(x, old_dist_spec, row_spec, pg)
      assert x.size(0) == 32 // size and x.size(1) == 32
      assert torch.cuda.memory_allocated() == orig_mem // size
      x.data = DistSpecManager._gather(x, row_spec, pg)
      assert torch.cuda.memory_allocated() == orig_mem
  
file_path: colossalai/tests/test_legacy/test_tensor/core/test_dist_spec_mgr.py
imports: import math
  import pytest
  import torch
  import torch.distributed
  import colossalai
  from colossalai.legacy.tensor import DistSpecManager, ProcessGroup, ReplicaSpec, ShardSpec
  from colossalai.testing import rerun_if_address_is_in_use, spawn
label_desc: Use image node_id 2 for calling a global function with example usage: check_mem() without return types