def METHOD_NAME(self) -> None:
    num_features = 1
    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 param in ebc.parameters():
        param._in_backward_optimizers = [
            WarmupOptimizer(
                param._in_backward_optimizers[0],
                [
                    WarmupStage(
                        policy=WarmupPolicy.LINEAR,
                        max_iters=10000,
                        value=0.5,
                        lr_scale=1.0,
                    )
                ],
                param_name="__warmup_state",
            )
        ]
        param._in_backward_optimizers[0].step()
        param._in_backward_optimizers[0].step()
        warmup_state = param._in_backward_optimizers[0].state_dict()["state"][
            "__warmup_state"
        ]
        _iter, _ = warmup_state["warmup"]
        self.assertEqual(_iter, 2)
        self.assertEqual(
            param._in_backward_optimizers[0].param_groups[0]["lr"], 0.05001
        )