    def test_schedules(self):
        """Make sure the math is correct."""
        for epochs, params, lr_checks, _ in self.cosine_schedule_cases:
            optimizer = self._get_optimizer()
            params["type"] = "cosine"
            scheduler = LearningRateScheduler.from_params(
                optimizer=optimizer, params=Params(params)
            )
            lrs = [optimizer.param_groups[0]["lr"]]
            for _ in range(epochs):
                scheduler.step()
                lrs.append(optimizer.param_groups[0]["lr"])

            for it, lr in lr_checks:
                assert lrs[it] == pytest.approx(lr), f"Iteration {it}: {lrs[it]} != {lr}"
Make sure the math is correct.
tests/training/learning_rate_schedulers/cosine_test.py