    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}"
Make sure scheduler will resume with the right state.
tests/training/learning_rate_schedulers/cosine_test.py