code: def test_custom_sharding_planner(self):
      megatron_lm = SimpleMegatronLM(
          [[17, 12], [12, 29]], rank=self.rank
      ).cuda(self.rank)
      planner = ChunkAllShardingPlanner(device_count=TEST_GPU_NUM)
      sharding_plan = planner.build_plan(megatron_lm)
  
      shard_module(megatron_lm, sharding_plan)
  
      # check to make sure the module already been sharded
      self.assertTrue(isinstance(megatron_lm.fc1.weight, ShardedTensor))
      self.assertTrue(isinstance(megatron_lm.fc2.weight, ShardedTensor))
      self.assertTrue(isinstance(megatron_lm.fc1.bias, ShardedTensor))
      self.assertTrue(isinstance(megatron_lm.fc2.bias, ShardedTensor))
  
imports: import sys
  import torch
  import torch.nn
  import torch.distributed
  from torch.testing._internal.common_distributed import requires_nccl, skip_if_lt_x_gpu
  from torch.distributed._shard import shard_module
  from torch.distributed._shard.sharding_plan import ShardingPlan, ShardingPlanner
  from torch.distributed._shard.sharding_spec import ChunkShardingSpec
  from torch.distributed._shard.sharded_tensor import ShardedTensor
  from torch.testing._internal.common_utils import TEST_WITH_DEV_DBG_ASAN, run_tests
  from torch.testing._internal.distributed._shard.sharded_tensor import TEST_GPU_NUM, ShardedTensorTestBase, with_comms
  from torch.testing._internal.distributed._shard.sharded_tensor._test_ops_common import generate_chunk_sharding_specs_for_test
  from torch.testing._internal.distributed._shard.test_common import SimpleMegatronLM
label_desc: Use image node_id 2 for calling the TestShardingPlan obj's underlying member method code with example usage: obj.test_custom_sharding_planner() without return types