def METHOD_NAME(self) -> None:
    num_features = 4
    tables = [
        EmbeddingBagConfig(
            num_embeddings=(i + 1) * 10,
            embedding_dim=(i + 1) * 4,
            name="table_" + str(i),
            feature_names=["feature_" + str(i)],
        )
        for i in range(num_features)
    ]
    ebc = EmbeddingBagCollection(tables=tables, device=torch.device("meta"))
    apply_optimizer_in_backward(
        RowWiseAdagrad,
        ebc.parameters(),
        {"lr": 0.01},
    )
    plan = construct_module_sharding_plan(
        ebc,
        apply_to_all(ebc, table_wise(rank=0)),
    )
    ebc = shard(
        module=ebc,
        plan=plan,
        device=self.curr_device,
    )
    for name, param in ebc.named_parameters():
        table_name = name[len("embedding_bags.") : -len("weight") - 1]
        self.assertEqual(
            param._in_backward_optimizers[0]
            .state_dict()["state"][""][f"{table_name}.momentum1"]
            .local_tensor()
            .data_ptr(),
            ebc._optim.state_dict()["state"][f"embedding_bags.{table_name}.weight"][
                f"{table_name}.momentum1"
            ]
            .local_tensor()
            .data_ptr(),
        )