code: def test_lr_scheduler_save_load():
      steps_per_checkpoint = 10
      total_steps = 100
      base_lr = 1.0
      reduce_limit = 3
  
      module = NumberInputFeature(
          NumberInputFeatureConfig(
              name="num1", encoder=DenseEncoderConfig()
          )
      )
      output1 = NumberOutputFeature(
          ECDNumberOutputFeatureConfig(name="output1", input_size=10),
          output_features={},
      )
  
      optimizer = SGD(module.parameters(), lr=base_lr)
      config = LRSchedulerConfig(
          warmup_fraction=0.2, reduce_on_plateau=reduce_limit
      )
      scheduler = LRScheduler(
          config=config,
          optimizer=optimizer,
          steps_per_checkpoint=steps_per_checkpoint,
          total_steps=total_steps,
      )
  
      progress_tracker = get_new_progress_tracker(
          batch_size=64,
          best_eval_metric_value=float("inf"),
          best_increase_batch_size_eval_metric=float("inf"),
          learning_rate=base_lr,
          output_features={"output1": output1},
      )
  
      for _ in range(10):
          scheduler.step()
  
      progress_tracker.train_metrics["output1"]["loss"].append(
          TrainerMetric(epoch=0, step=10, value=1.0)
      )
      scheduler.eval_step(progress_tracker, "output1")
  
      optimizer_state = optimizer.state_dict()
      scheduler_state = scheduler.state_dict()
  
      optimizer2 = SGD(module.parameters(), lr=base_lr)
      scheduler2 = LRScheduler(
          config=config,
          optimizer=optimizer2,
          steps_per_checkpoint=steps_per_checkpoint,
          total_steps=total_steps,
      )
  
      # Important: state needs to be loaded after init of optimizer and scheduler, otherwise
      # it can override loaded state
      optimizer2.load_state_dict(optimizer_state)
      scheduler2.load_state_dict(scheduler_state)
  
      lr = optimizer.param_groups[0]["lr"]
      assert lr == optimizer2.param_groups[0]["lr"]
      assert scheduler.state_dict() == scheduler2.state_dict()
  
      for _ in range(10):
          scheduler.step()
          scheduler2.step()
  
      progress_tracker.train_metrics["output1"]["loss"].append(
          TrainerMetric(epoch=1, step=20, value=0.8)
      )
      scheduler.eval_step(progress_tracker, "output1")
      scheduler2.eval_step(progress_tracker, "output1")
  
      assert lr != optimizer.param_groups[0]["lr"]
      assert (
          optimizer.param_groups[0]["lr"]
          == optimizer2.param_groups[0]["lr"]
      )
      assert scheduler.state_dict() == scheduler2.state_dict()
  
file_path: ludwig/tests/ludwig/modules/test_lr_scheduler.py
imports: import math
  import numpy
  from torch.optim import SGD
  from ludwig.features.number_feature import NumberInputFeature, NumberOutputFeature
  from ludwig.modules.lr_scheduler import LRScheduler
  from ludwig.schema.encoders.base import DenseEncoderConfig
  from ludwig.schema.features.number_feature import ECDNumberOutputFeatureConfig, NumberInputFeatureConfig
  from ludwig.schema.lr_scheduler import LRSchedulerConfig
  from ludwig.utils.metric_utils import TrainerMetric
  from ludwig.utils.trainer_utils import get_new_progress_tracker
label_desc: Use image node_id 5 for calling a global function with example usage: test_lr_scheduler_save_load() without return types