Regression in loading optimizer learning rate about pytorch HOT 4 OPEN

def _init_state_dict(state_dict: STATE_DICT_TYPE) -> None:
    state_dict_assigned_storage = tree_map_only(
        torch.Tensor, lambda v: _init_meta_tensor(v), state_dict
    )
    # The inplace version of tree_map_only, tree_map_only_ doesn't seem to work.
    # So we need to temporariy update the each element in the state dict with meta tensor.
    for k in state_dict.keys():
        state_dict[k] = state_dict_assigned_storage[k]

def _init_state_dict(state_dict: STATE_DICT_TYPE) -> None:
    tree_map_only_(torch.Tensor, _init_meta_tensor, state_dict)

def load_checkpoint(checkpoint_path: str, model: torch.nn.Module, optimizer: torch.optim.Optimizer, scheduler: torch.optim.lr_scheduler.LambdaLR) -> dict:
   """Load the model and optimizer state from a checkpoint folder
   
   Args:
       checkpoint_path: the path to the checkpoint folder
       model: the model to load
       optimizer: the optimizer to load
       scheduler: the scheduler to load
   """
   # 1. Load distributed states
   fs_storage_reader = dcp.FileSystemReader(checkpoint_path)

   model_state_dict, optimizer_state_dict = get_state_dict(model, optimizer)
   dcp_state_dict = {
       "model": model_state_dict,
       "optimizer": optimizer_state_dict,
   }
   dcp.load(dcp_state_dict, storage_reader=fs_storage_reader)
   assert optimizer_state_dict["param_groups"][0]["lr"] == dcp_state_dict["optimizer"]["param_groups"][0]["lr"]
   set_state_dict(model, optimizer, model_state_dict=dcp_state_dict["model"], optim_state_dict=dcp_state_dict["optimizer"])
   
   # 2. Load global states
   global_state_dict = torch.load(os.path.join(checkpoint_path, "global_state_dict.pt"))
   scheduler.load_state_dict(global_state_dict["scheduler"])