code: def test_set_learning_rate(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,
          )
  
file_path: torchrec/torchrec/distributed/composable/tests/test_fused_optim.py
imports: import os
  import unittest
  import torch
  from torch import distributed
  from torch.distributed.optim import _apply_optimizer_in_backward
  from torchrec.distributed.shard import shard
  from torchrec.distributed.sharding_plan import apply_to_all, construct_module_sharding_plan, table_wise
  from torchrec.modules.embedding_configs import EmbeddingBagConfig
  from torchrec.modules.embedding_modules import EmbeddingBagCollection
  from torchrec.optim.rowwise_adagrad import RowWiseAdagrad
  from torchrec.optim.warmup import WarmupOptimizer, WarmupPolicy, WarmupStage
  from torchrec.test_utils import get_free_port
label_desc: Use image node_id 4 for calling the TestFusedOptim obj's underlying member method code with example usage: obj.test_set_learning_rate() without return types