def METHOD_NAME(self):
    rowwise_spec = ChunkShardingSpec(
        dim=0,
        placements=[
            "rank:0/cuda:0",
            "rank:1/cuda:1",
            "rank:2/cuda:2",
            "rank:3/cuda:3",
        ],
    )
    sharded_model = MyShardedModel(spec=rowwise_spec).cuda()
    sharded_model_params = dict(sharded_model.named_parameters())
    param_keys = list(sharded_model_params.keys())
    self.assertEqual(len(param_keys), 2)
    self.assertTrue("param" in param_keys)
    self.assertTrue("sharded_param" in param_keys)
    sharded_linear = MyShardedLinear(rank=self.rank).cuda()
    sharded_linear.shard_parameter()
    sharded_linear_params = dict(sharded_linear.named_parameters())
    param_keys = list(sharded_linear_params.keys())
    self.assertEqual(len(param_keys), 4)
    self.assertTrue("linear1.bias" in param_keys)
    self.assertTrue("linear2.bias" in param_keys)
    self.assertTrue("linear1.weight" in param_keys)
    self.assertTrue("linear2.weight" in param_keys)
    self.assertFalse("bias" in param_keys)