def METHOD_NAME(self):
    class MyModule(nn.Module):
        def __init__(self):
            super().__init__()
            self.ebc = CustomEmbeddingBagCollection(10, 10, 8)
        def forward(self, inputs):
            return self.ebc(inputs)
    custom_sharder = CustomSharder(
        devices=[f"rank:{i}/cuda:{i}" for i in range(TEST_GPU_NUM)],
        split_sharding_idx=TEST_GPU_NUM // 2
    )
    sharding_plan = ShardingPlan(
        plan={
            "ebc": custom_sharder,
        })
    local_model = MyModule().cuda(self.rank)
    sharded_model = copy.deepcopy(local_model)
    # shard the module with the provided sharding plan
    shard_module(sharded_model, sharding_plan)
    # check to make sure the module already been sharded
    emb_bags = sharded_model.ebc.embedding_bags
    self.assertTrue(isinstance(emb_bags["embedding_bag_0"].weight, ShardedTensor))
    self.assertTrue(isinstance(emb_bags["embedding_bag_9"].weight, ShardedTensor))
    self.assertEqual(emb_bags["embedding_bag_0"].weight.sharding_spec(), custom_sharder.rowwise_spec)
    self.assertEqual(emb_bags["embedding_bag_9"].weight.sharding_spec(), custom_sharder.colwise_spec)
    # make sure we can run sharded computation and compare outputs
    # with the local model version
    input = torch.arange(8).reshape((2, 4)).cuda(self.rank)
    local_output = local_model(input)
    sharded_output = sharded_model(input)
    self.assertEqual(local_output, sharded_output)