Repository: pytorch/ignite

Issue:
Make LRScheduler attachable to Events.ITERATION_STARTED
## 🚀 Feature

Currently, the correct way to use `LRScheduler` wrapper for pytorch >= 1.1.0 is the following:
```python
from torch.optim.lr_scheduler import StepLR

torch_lr_scheduler = StepLR(optimizer, step_size=3, gamma=0.1)
scheduler = LRScheduler(torch_lr_scheduler)

@trainer.on(Events.ITERATION_COMPLETED)
def print_lr():
    print(optimizer.param_groups[0]["lr"])

# In this example, we assume to have installed PyTorch>=1.1.0
# (with new `torch.optim.lr_scheduler` behaviour) and
# we attach scheduler to Events.ITERATION_COMPLETED
# instead of Events.ITERATION_STARTED to make sure to use
# the first lr value from the optimizer, otherwise it is will be skipped:
trainer.add_event_handler(Events.ITERATION_COMPLETED, scheduler)

trainer..run([0] * 8, max_epochs=1)
```
>
```
0.1
0.1
0.1
0.010
0.010
0.010
0.001
0.001
```

- https://github.com/pytorch/ignite/pull/2463

however, other schedulers should be used as following ([link](https://pytorch.org/ignite/generated/ignite.handlers.param_scheduler.PiecewiseLinear.html#ignite.handlers.param_scheduler.PiecewiseLinear))
```python
milestones_values = [(1, 1.0), (3, 0.8), (5, 0.2)]
scheduler = PiecewiseLinear(
    optimizer, "lr", milestones_values=milestones_values)

trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)

@trainer.on(Events.ITERATION_COMPLETED)
def print_lr():
    print(optimizer.param_groups[0]["lr"])

trainer.run([0] * 6, max_epochs=1)
```

The idea is to improve `LRScheduler` such that we could attach it to `Events.ITERATION_STARTED` and have a coherent API. It will be a BC-breaking change, but for good.

So, desired example using `LRScheduler` should be:
```python
from torch.optim.lr_scheduler import StepLR

torch_lr_scheduler = StepLR(optimizer, step_size=3, gamma=0.1)
scheduler = LRScheduler(torch_lr_scheduler)

trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)

@trainer.on(Events.ITERATION_COMPLETED)
def print_lr():
    print(optimizer.param_groups[0]["lr"])

trainer.run([0] * 8, max_epochs=1)
```
Currently, this gives a wrong behaviour as the first 0.1 wasn't consumed by the training step.
```
0.1
0.1
0.010
0.010
0.010
0.001
0.001
```
The idea could be to retain the first value and reapply it once and then keep everything as it is now.




Discussion:
sdesrozis: @yuta0821 Thanks very much for this work. I have to think about the api though. The `use_legacy` arg is for BC (so to be removed one day) but the internal only `keep_fist_lr` makes me a little doubtful. This arg is only required using ConcatScheduler, could we do smarter ? Anyway, it works and that is pretty cool. Let's just think about it. I suppose that @vfdev-5 would have an opinion on that.
yuta0821: @sdesrozis Yes, I agree, I don't want to add more args without any thought. However, I haven't yet figured out how to control the `lr` value of an `LRScheduler` bound to a `ConcatScheduler` without using `keep_first_lr`. Originally, this was implemented by adjusting the `milestones_values` of `PeacewiseLinear`,  but chage of this PR no longer works. I'll see if there's a smarter way to do this.

sdesrozis: My toughts are the target would be having the following code working fine
```python
from torch.optim.lr_scheduler import StepLR

scheduler1 = StepLR(optimizer, step_size=3, gamma=0.1)
scheduler1 = LRScheduler(scheduler1)

scheduler2 = StepLR(optimizer, step_size=3, gamma=0.01)
scheduler2 = LRScheduler(scheduler2)

scheduler = ConcatScheduler(schedulers=[scheduler1, scheduler2], durations=[4, ])
```

Whether `use_legacy` is used or not for `scheduler2`, it should produce the same output. It means `ConcatScheduler` should enable `keep_first_lr` of `LRScheduler` internally without any user api. The minimal would be to enable using a method rather than at init.
```python
scheduler2 = LRScheduler(scheduler2)
scheduler2.keep_first_lr = True
```
I don't know yet...
yuta0821: @sdesrozis I'd like to confirm the behavior of following code using `use_legacy=True`. Is this output is right ?
```python
from torch.optim.lr_scheduler import StepLR

optimizer = torch.optim.SGD([tensor], lr=0.01)
torch_scheduler1 = StepLR(optimizer, step_size=3, gamma=0.1)
scheduler1 = LRScheduler(torch_scheduler1, use_legacy=True)

torch_scheduler2 = StepLR(optimizer, step_size=3, gamma=0.01)
scheduler2 = LRScheduler(torch_scheduler2, use_legacy=True)

scheduler = ConcatScheduler(schedulers=[scheduler1, scheduler2], durations=[4, ])

def dummy_update(engine, batch):
    pass
trainer = Engine(dummy_update)

@trainer.on(Events.ITERATION_COMPLETED)
def print_lr():
    print(optimizer.param_groups[0]["lr"])
trainer.add_event_handler(Events.ITERATION_COMPLETED, scheduler)

_ = trainer.run([0] * 9, max_epochs=1)
```
```
0.01
0.01
0.01
0.001
0.001
0.001
0.001
1e-05
1e-05
```
> Whether `use_legacy` is used or not for schedulers, it should produce the same output. It means `ConcatScheduler` should enable `keep_first_lr` of `LRScheduler` internally without any user api.

OK, I understand what message means. Thank you for the clear explanation. Surely, we must enable `keep_first_lr` of `LRScheduler` internally without any user api.
vfdev-5: > @sdesrozis I'd like to confirm the behavior of following code using `use_legacy=True`. Is this output is right ?
> 
> ```python
> from torch.optim.lr_scheduler import StepLR
> 
> optimizer = torch.optim.SGD([tensor], lr=0.01)
> torch_scheduler1 = StepLR(optimizer, step_size=3, gamma=0.1)
> scheduler1 = LRScheduler(torch_scheduler1, use_legacy=True)
> 
> torch_scheduler2 = StepLR(optimizer, step_size=3, gamma=0.01)
> scheduler2 = LRScheduler(torch_scheduler2, use_legacy=True)
> 
> scheduler = ConcatScheduler(schedulers=[scheduler1, scheduler2], durations=[4, ])
> 
> def dummy_update(engine, batch):
>     pass
> trainer = Engine(dummy_update)
> 
> @trainer.on(Events.ITERATION_COMPLETED)
> def print_lr():
>     print(optimizer.param_groups[0]["lr"])
> trainer.add_event_handler(Events.ITERATION_COMPLETED, scheduler)
> 
> _ = trainer.run([0] * 9, max_epochs=1)
> ```
> 
> ```
> 0.01
> 0.01
> 0.01
> 0.001
> 0.001
> 0.001
> 0.001
> 1e-05
> 1e-05
> ```

@yuta0821 the output looks correct to me. We have 
```
--- scheduler 1 
0.01
0.01
0.01
0.001
--- scheduler 2
0.001
0.001
0.001
1e-05
1e-05
```

What happens if using `use_legacy=False` and attaching to appropriate event ?
yuta0821: @vfdev-5 
Sorry for late reply.
> What happens if using use_legacy=False and attaching to appropriate event ?

If set `trainer.add_event_handler` to the proper position, the behavior with `use_legacy = False` becomes as follows.
With `keep_lr_first`, the behavior of this code is the same as with `use_legacy = True`.

```python
from torch.optim.lr_scheduler import StepLR

optimizer = torch.optim.SGD([tensor], lr=0.01)
torch_scheduler1 = StepLR(optimizer, step_size=3, gamma=0.1)
scheduler1 = LRScheduler(torch_scheduler1, use_legacy=False)

torch_scheduler2 = StepLR(optimizer, step_size=3, gamma=0.01)
scheduler2 = LRScheduler(torch_scheduler2, use_legacy=False)

scheduler = ConcatScheduler(schedulers=[scheduler1, scheduler2], durations=[4, ])

def dummy_update(engine, batch):
    pass
trainer = Engine(dummy_update)
trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)

@trainer.on(Events.ITERATION_COMPLETED)
def print_lr():
    print(optimizer.param_groups[0]["lr"])
    
_ = trainer.run([0] * 9, max_epochs=1)
```
```
--- scheduler 1 
0.01
0.01
0.01
0.001
--- scheduler 2 
0.001
0.001
0.001
1e-05
1e-05
```
vfdev-5: @yuta0821 sorry for delay on this PR, i'm still missing a complete understanding of the problem here. I'll try to figure out what happens exactly on my side and try to suggest a way to handle that. Adding an argument for internal usage is not a good design.
yuta0821: @vfdev-5 

> @yuta0821 sorry for delay on this PR, i'm still missing a complete understanding of the problem here. I'll try to figure out what happens exactly on my side and try to suggest a way to handle that. Adding an argument for internal usage is not a good design.
> Thanks for your patience and working on the PR !

No, sir, not at all. Thanks a lot for your help. If I can contribute in any way, I'll be happy to do so.
vfdev-5: @yuta0821 i understand the problem and why you need to introduce `keep_first_lr` argument.
In two words for others, the problem with `create_lr_scheduler_with_warmup` is that either 1) we have an additional transitional lr value between `PiecewiseLinear` and `LRScheduler`:
```- PiecewiseLinear updates LR
1 0.0
- PiecewiseLinear updates LR
2 0.03
- PiecewiseLinear updates LR
3 0.06
- PiecewiseLinear updates LR
4 0.09
- PiecewiseLinear updates LR
5 0.12
- PiecewiseLinear updates LR  <----- additional step to remove
6 0.01
- LRScheduler updates LR 
7 0.01
- LRScheduler updates LR
8 0.01
- LRScheduler updates LR
9 0.01
- LRScheduler updates LR
10 0.005
```
or 2) `LRScheduler` modifies LR in a relative way => if `PiecewiseLinear` stops one step before, then `LRScheduler` will work from this value, e.g `0.12` from example above. This also gives a wrong behavior.

I see here two ways to move forward:
- either we add a more meaningful public arg to `LRScheduler`, e.g. `relative: bool = True` and thus we setup `create_lr_scheduler_with_warmup` with `LRScheduler(..., relative=False)`
- do something like `lr_scheduler.lr_scheduler.last_epoch += 1` inside `create_lr_scheduler_with_warmup`:
```python
lr_scheduler = LRScheduler(lr_scheduler, save_history=save_history)
lr_scheduler.lr_scheduler.last_epoch += 1
```
yuta0821: > @yuta0821 i understand the problem and why you need to introduce keep_first_lr argument.

Thank you for your analysis. Your summary is exactly what I was thinking.

vfdev-5: @yuta0821 I'm trying to figure out how to merge your work to master step by step.
You updated tests/ignite/handlers/test_param_scheduler.py , do you think we could make merge without updates introduced in ignite/handlers/param_scheduler.py ?
If yes, could you please send a PR with only updated tests ? Thanks a lot for your patience
yuta0821: I think we could make merge neither `tests/ignite/handlers/test_param_scheduler.py` nor `tests/ignite/contrib/engines/test_common.py`  without updates introduced `ignite/handlers/param_scheduler.py`.
This is because the changes in tests such like `trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)` are added to tackle on changes caused by `ignite/handlers/param_scheduler.py`.
vfdev-5: @yuta0821 I pushed few commits to make this PR merged. Basically, it is one of your suggestions. Currently, I do not have a better way to fix it. Let's keep it as it is.
Thanks again for starting this work and sorry for delay it took to make it landed

Patch:
diff --git a/ignite/contrib/engines/common.py b/ignite/contrib/engines/common.py
index 3b773f11..bac1eaa5 100644
--- a/ignite/contrib/engines/common.py
+++ b/ignite/contrib/engines/common.py
@@ -13,7 +13,6 @@ import ignite.distributed as idist
 from ignite.contrib.handlers import (
     ClearMLLogger,
     global_step_from_engine,
-    LRScheduler,
     MLflowLogger,
     NeptuneLogger,
     PolyaxonLogger,
@@ -165,8 +164,6 @@ def _setup_common_training_handlers(
             trainer.add_event_handler(
                 Events.ITERATION_COMPLETED, lambda engine: cast(_LRScheduler, lr_scheduler).step()
             )
-        elif isinstance(lr_scheduler, LRScheduler):
-            trainer.add_event_handler(Events.ITERATION_COMPLETED, lr_scheduler)
         else:
             trainer.add_event_handler(Events.ITERATION_STARTED, lr_scheduler)
 
diff --git a/ignite/handlers/param_scheduler.py b/ignite/handlers/param_scheduler.py
index e88c24a3..b09611c7 100644
--- a/ignite/handlers/param_scheduler.py
+++ b/ignite/handlers/param_scheduler.py
@@ -795,6 +795,7 @@ class LRScheduler(ParamScheduler):
         lr_scheduler: lr_scheduler object to wrap.
         save_history: whether to log the parameter values to
             `engine.state.param_history`, (default=False).
+        use_legacy: if True, scheduler should be attached to ``Events.ITERATION_COMPLETED``, (default=False).
 
     Examples:
 
@@ -808,20 +809,14 @@ class LRScheduler(ParamScheduler):
             from torch.optim.lr_scheduler import StepLR
 
             torch_lr_scheduler = StepLR(default_optimizer, step_size=3, gamma=0.1)
-
             scheduler = LRScheduler(torch_lr_scheduler)
 
+            default_trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
+
             @default_trainer.on(Events.ITERATION_COMPLETED)
             def print_lr():
                 print(default_optimizer.param_groups[0]["lr"])
 
-            # In this example, we assume to have installed PyTorch>=1.1.0
-            # (with new `torch.optim.lr_scheduler` behaviour) and
-            # we attach scheduler to Events.ITERATION_COMPLETED
-            # instead of Events.ITERATION_STARTED to make sure to use
-            # the first lr value from the optimizer, otherwise it is will be skipped:
-            default_trainer.add_event_handler(Events.ITERATION_COMPLETED, scheduler)
-
             default_trainer.run([0] * 8, max_epochs=1)
 
         .. testoutput::
@@ -836,9 +831,17 @@ class LRScheduler(ParamScheduler):
             0.001...
 
     .. versionadded:: 0.4.5
+
+    ..  versionchanged:: 0.5.0
+        added `use_legacy` argument
     """
 
-    def __init__(self, lr_scheduler: _LRScheduler, save_history: bool = False):
+    def __init__(
+        self,
+        lr_scheduler: _LRScheduler,
+        save_history: bool = False,
+        use_legacy: bool = False,
+    ):
 
         if not isinstance(lr_scheduler, _LRScheduler):
             raise TypeError(
@@ -852,11 +855,19 @@ class LRScheduler(ParamScheduler):
             param_name="lr",
             save_history=save_history,
         )
+        if use_legacy:
+            warnings.warn(
+                "Please make sure to attach scheduler to Events.ITERATION_COMPLETED "
+                "instead of Events.ITERATION_STARTED to make sure to use "
+                "the first lr value from the optimizer, otherwise it is will be skipped"
+            )
+            self.lr_scheduler.last_epoch += 1  # type: ignore[attr-defined]
+
         self._state_attrs += ["lr_scheduler"]
 
     def __call__(self, engine: Optional[Engine], name: Optional[str] = None) -> None:
-        self.lr_scheduler.last_epoch += 1  # type: ignore[attr-defined]
         super(LRScheduler, self).__call__(engine, name)
+        self.lr_scheduler.last_epoch += 1  # type: ignore[attr-defined]
 
     def get_param(self) -> Union[float, List[float]]:
         """Method to get current optimizer's parameter value"""
@@ -904,9 +915,9 @@ class LRScheduler(ParamScheduler):
             values = []
             scheduler = cls(save_history=False, lr_scheduler=lr_scheduler, **kwargs)
             for i in range(num_events):
+                scheduler(engine=None)
                 params = [p[scheduler.param_name] for p in scheduler.optimizer_param_groups]
                 values.append([i] + params)
-                scheduler(engine=None)
 
             obj = torch.load(cache_filepath.as_posix())
             lr_scheduler.load_state_dict(obj["lr_scheduler"])
@@ -927,8 +938,7 @@ def create_lr_scheduler_with_warmup(
     Helper method to create a learning rate scheduler with a linear warm-up.
 
     Args:
-        lr_scheduler: learning rate scheduler
-            after the warm-up.
+        lr_scheduler: learning rate scheduler after the warm-up.
         warmup_start_value: learning rate start value of the warm-up phase.
         warmup_duration: warm-up phase duration, number of events.
         warmup_end_value: learning rate end value of the warm-up phase, (default=None). If None,
@@ -1011,10 +1021,15 @@ def create_lr_scheduler_with_warmup(
 
         if isinstance(lr_scheduler, _LRScheduler):
             init_lr = param_group["lr"]
-
             if init_lr != param_group_warmup_end_value:
                 milestones_values.append((warmup_duration, init_lr))
 
+            # We need to advance torch lr_scheduler to avoid duplicated lr value
+            # given by PiecewiseLinear and LRScheduler.
+            # We suggest to attach output scheduler on ITERATION_STARTED but
+            # torch lr_scheduler works with ITERATION_COMPLETED
+            # See also https://github.com/pytorch/ignite/pull/2496#issuecomment-1065984440
+            lr_scheduler.last_epoch += 1
             lr_scheduler = LRScheduler(lr_scheduler, save_history=save_history)
         else:
             init_lr = lr_scheduler.get_param()


Test patch:
diff --git a/tests/ignite/contrib/engines/test_common.py b/tests/ignite/contrib/engines/test_common.py
index 9e89491a..cadae338 100644
--- a/tests/ignite/contrib/engines/test_common.py
+++ b/tests/ignite/contrib/engines/test_common.py
@@ -128,8 +128,8 @@ def _test_setup_common_training_handlers(
 
     # Check LR scheduling
     assert optimizer.param_groups[0]["lr"] <= lr * gamma ** (
-        num_iters * num_epochs / step_size
-    ), f"{optimizer.param_groups[0]['lr']} vs {lr * gamma ** (num_iters * num_epochs / step_size)}"
+        (num_iters * num_epochs - 1) // step_size
+    ), f"{optimizer.param_groups[0]['lr']} vs {lr * gamma ** ((num_iters * num_epochs - 1) // step_size)}"
 
 
 def test_asserts_setup_common_training_handlers():
diff --git a/tests/ignite/handlers/test_param_scheduler.py b/tests/ignite/handlers/test_param_scheduler.py
index 77a9a3eb..8edd0225 100644
--- a/tests/ignite/handlers/test_param_scheduler.py
+++ b/tests/ignite/handlers/test_param_scheduler.py
@@ -648,53 +648,71 @@ def test_lr_scheduler(torch_lr_scheduler_cls, kwargs):
     tensor = torch.zeros([1], requires_grad=True)
     optimizer1 = torch.optim.SGD([tensor], lr=0.01)
     optimizer2 = torch.optim.SGD([tensor], lr=0.01)
+    optimizer3 = torch.optim.SGD([tensor], lr=0.01)
     opt_state_dict1 = optimizer1.state_dict()
     opt_state_dict2 = optimizer2.state_dict()
+    opt_state_dict3 = optimizer3.state_dict()
 
     torch_lr_scheduler1 = torch_lr_scheduler_cls(optimizer=optimizer1, **kwargs)
-    scheduler = LRScheduler(torch_lr_scheduler1)
-    state_dict1 = scheduler.state_dict()
+    scheduler1 = LRScheduler(torch_lr_scheduler1)
+    state_dict1 = scheduler1.state_dict()
 
     torch_lr_scheduler2 = torch_lr_scheduler_cls(optimizer=optimizer2, **kwargs)
-    state_dict2 = torch_lr_scheduler2.state_dict()
+    with pytest.warns(UserWarning, match=r"the first lr value from the optimizer, otherwise it is will be skipped"):
+        scheduler2 = LRScheduler(torch_lr_scheduler2, use_legacy=True)
+    state_dict2 = scheduler2.state_dict()
+
+    torch_lr_scheduler3 = torch_lr_scheduler_cls(optimizer=optimizer3, **kwargs)
+    state_dict3 = torch_lr_scheduler3.state_dict()
 
     def dummy_update(engine, batch):
         optimizer1.step()
         optimizer2.step()
+        optimizer3.step()
 
     trainer = Engine(dummy_update)
+    trainer.add_event_handler(Events.ITERATION_STARTED, scheduler1)
 
     @trainer.on(Events.ITERATION_STARTED)
-    def save_lr(engine):
-        lrs.append(optimizer1.param_groups[0]["lr"])
+    def save_lr1(engine):
+        lrs1.append(optimizer1.param_groups[0]["lr"])
+
+    @trainer.on(Events.ITERATION_STARTED)
+    def save_lr2(engine):
+        lrs2.append(optimizer2.param_groups[0]["lr"])
 
     @trainer.on(Events.ITERATION_STARTED)
     def save_true_lr(engine):
-        lrs_true.append(optimizer2.param_groups[0]["lr"])
+        lrs_true.append(optimizer3.param_groups[0]["lr"])
 
     @trainer.on(Events.ITERATION_COMPLETED)
     def torch_lr_scheduler_step(engine):
-        torch_lr_scheduler2.step()
+        torch_lr_scheduler3.step()
 
-    trainer.add_event_handler(Events.ITERATION_COMPLETED, scheduler)
+    trainer.add_event_handler(Events.ITERATION_COMPLETED, scheduler2)
 
     for _ in range(2):
-        lrs = []
+        lrs1 = []
+        lrs2 = []
         lrs_true = []
         data = [0] * 10
         max_epochs = 2
         trainer.run(data, max_epochs=max_epochs)
-        assert lrs_true == pytest.approx(lrs), f"{_}: {lrs_true} ({len(lrs_true)}) vs {lrs} ({len(lrs)})"
+        assert lrs_true == pytest.approx(lrs1), f"{_}: {lrs_true} ({len(lrs_true)}) vs {lrs1} ({len(lrs1)})"
+        assert lrs_true == pytest.approx(lrs2), f"{_}: {lrs_true} ({len(lrs_true)}) vs {lrs2} ({len(lrs2)})"
         optimizer1.load_state_dict(opt_state_dict1)
-        scheduler.load_state_dict(state_dict1)
+        scheduler1.load_state_dict(state_dict1)
         optimizer2.load_state_dict(opt_state_dict2)
-        torch_lr_scheduler2.load_state_dict(state_dict2)
+        scheduler2.load_state_dict(state_dict2)
+        optimizer3.load_state_dict(opt_state_dict3)
+        torch_lr_scheduler3.load_state_dict(state_dict3)
 
-    optimizer3 = torch.optim.SGD([tensor], lr=0.01)
-    torch_lr_scheduler3 = torch_lr_scheduler_cls(optimizer=optimizer3, **kwargs)
+    optimizer4 = torch.optim.SGD([tensor], lr=0.01)
+    torch_lr_scheduler4 = torch_lr_scheduler_cls(optimizer=optimizer4, **kwargs)
 
-    simulated_values = LRScheduler.simulate_values(num_events=len(data) * max_epochs, lr_scheduler=torch_lr_scheduler3)
-    assert lrs == pytest.approx([v for i, v in simulated_values])
+    simulated_values = LRScheduler.simulate_values(num_events=len(data) * max_epochs, lr_scheduler=torch_lr_scheduler4)
+    assert lrs1 == pytest.approx([v for i, v in simulated_values])
+    assert lrs2 == pytest.approx([v for i, v in simulated_values])
 
 
 def test_piecewiselinear_asserts():
@@ -813,11 +831,8 @@ def test_simulate_and_plot_values():
 
     def _test(scheduler_cls, **scheduler_kwargs):
 
-        optimizer = None
-        event = Events.ITERATION_STARTED
         if scheduler_cls == LRScheduler:
             optimizer = scheduler_kwargs["lr_scheduler"].optimizer
-            event = Events.ITERATION_COMPLETED
         elif scheduler_cls == ConcatScheduler:
             optimizer = scheduler_kwargs["optimizer"]
             del scheduler_kwargs["optimizer"]
@@ -828,7 +843,7 @@ def test_simulate_and_plot_values():
 
         max_epochs = 2
         data = [0] * 10
-        # simulated_values = scheduler_cls.simulate_values(num_events=len(data) * max_epochs, **scheduler_kwargs)
+        simulated_values = scheduler_cls.simulate_values(num_events=len(data) * max_epochs, **scheduler_kwargs)
 
         scheduler = scheduler_cls(**scheduler_kwargs)
 
@@ -838,15 +853,11 @@ def test_simulate_and_plot_values():
             lrs.append(optimizer.param_groups[0]["lr"])
 
         trainer = Engine(lambda engine, batch: None)
-        trainer.add_event_handler(event, scheduler)
+        trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
         trainer.add_event_handler(Events.ITERATION_STARTED, save_lr)
         trainer.run(data, max_epochs=max_epochs)
 
-        # assert lrs == pytest.approx([v for i, v in simulated_values])
-
-        if scheduler_cls == LRScheduler or scheduler_cls == ConcatScheduler:
-            # As internal state of torch lr scheduler has been changed the following checks will fail
-            return
+        assert lrs == pytest.approx([v for i, v in simulated_values])
 
         # reexecute to check if no internal changes
         # simulated_values = scheduler_cls.simulate_values(num_events=len(data) * max_epochs,
@@ -937,7 +948,7 @@ def test_create_lr_scheduler_with_warmup_asserts():
 
 
 @pytest.mark.parametrize(
-    "lrsched_warmup_config",
+    "lr_scheduler_name, warmup_start_value, warmup_end_value, warmup_duration, warmup_end_next_value",
     [
         # A) opt lr != warmup_end_value
         ("ExponentialLR", 0.01, 0.05, 10, 0.2),
@@ -955,15 +966,9 @@ def test_create_lr_scheduler_with_warmup_asserts():
         ("ExponentialLR", 0.01, None, 10, 0.2 * 0.98),
     ],
 )
-def test_create_lr_scheduler_with_warmup(lrsched_warmup_config):
-
-    (
-        lr_scheduler_name,
-        warmup_start_value,
-        warmup_end_value,
-        warmup_duration,
-        warmup_end_next_value,
-    ) = lrsched_warmup_config
+def test_create_lr_scheduler_with_warmup(
+    lr_scheduler_name, warmup_start_value, warmup_end_value, warmup_duration, warmup_end_next_value
+):
 
     t1 = torch.zeros([1], requires_grad=True)
 
@@ -981,6 +986,11 @@ def test_create_lr_scheduler_with_warmup(lrsched_warmup_config):
     num_iterations = 10
     max_epochs = 20
 
+    if warmup_end_value is None:
+        expected_warmup_end_value = optimizer.param_groups[0]["lr"]
+    else:
+        expected_warmup_end_value = warmup_end_value
+
     simulated_values = [None] * (num_iterations * max_epochs)
     scheduler = create_lr_scheduler_with_warmup(
         lr_scheduler,
@@ -989,8 +999,6 @@ def test_create_lr_scheduler_with_warmup(lrsched_warmup_config):
         warmup_duration=warmup_duration,
         output_simulated_values=simulated_values,
     )
-    if warmup_end_value is None:
-        warmup_end_value = optimizer.param_groups[0]["lr"]
 
     state_dict = scheduler.state_dict()
     trainer = Engine(lambda engine, batch: None)
@@ -1007,11 +1015,11 @@ def test_create_lr_scheduler_with_warmup(lrsched_warmup_config):
         lrs = []
         trainer.run(data, max_epochs=max_epochs)
 
-        assert lrs == pytest.approx([v for i, v in simulated_values])
+        assert lrs == pytest.approx([v for _, v in simulated_values])
 
         assert lrs[0] == pytest.approx(warmup_start_value), f"lrs={lrs[: warmup_duration + num_iterations]}"
         assert lrs[warmup_duration - 1] == pytest.approx(
-            warmup_end_value
+            expected_warmup_end_value
         ), f"lrs={lrs[: warmup_duration + num_iterations]}"
         assert lrs[warmup_duration] == pytest.approx(
             warmup_end_next_value