docstr: Make sure scheduler will resume with the right state.
docstr_tok: ["Make","sure","scheduler","will","resume","with","the","right","state","."]
code: def test_schedules_with_save_and_resume(self):
      """Make sure scheduler will resume with the right state."""
  
      def init_and_restore_scheduler(
          optimizer: torch.optim.Optimizer,
          params: Dict[str, Any],
          state_dict: Dict[str, Any] = None,
      ):
          """
          Initialize a new scheduler and optionally restore its state from
          a checkpoint.
          """
          params["type"] = "cosine"
          scheduler = LearningRateScheduler.from_params(
              optimizer=optimizer, params=Params(deepcopy(params))
          )
          if state_dict is not None:
              scheduler.load_state_dict(state_dict)
          return scheduler
  
      for (
          epochs,
          params,
          lr_checks,
          checkpoints,
      ) in self.cosine_schedule_cases:
          optimizer = self._get_optimizer()
          scheduler = init_and_restore_scheduler(optimizer, params)
          state = scheduler.state_dict()
  
          lrs = [optimizer.param_groups[0]["lr"]]
          for epoch in range(epochs):
              if epoch in checkpoints:
                  # Restore scheduler from state dict.
                  scheduler = init_and_restore_scheduler(
                      optimizer, params, state_dict=state
                  )
  
              # Take step and record learning rate.
              scheduler.step(1)
              lrs.append(optimizer.param_groups[0]["lr"])
  
              # Save state again.
              state = scheduler.state_dict()
  
          for it, lr in lr_checks:
              assert lrs[it] == pytest.approx(
                  lr
              ), f"Iteration {it}: {lrs[it]} != {lr}"
  
file_path: allennlp/tests/training/learning_rate_schedulers/cosine_test.py
imports: from copy import deepcopy
  from typing import Dict, Any
  import torch
  import pytest
  from allennlp.common import Params
  from allennlp.common.checks import ConfigurationError
  from allennlp.common.testing import AllenNlpTestCase
  from allennlp.training.learning_rate_schedulers import LearningRateScheduler
  from allennlp.training.optimizers import Optimizer
label_desc: Use image node_id 5 for calling the CosineWithRestartsTest obj's underlying member method code with example usage: obj.test_schedules_with_save_and_resume() and returns: scheduler