Repository: pytorch/ignite

Issue:
How to find suitable lrs with FastaiLRFinder when the optimizer has Multiple groups ?
```python
optimizer = optim.SGD([
    {'params': model.conv.parameters(), 'lr': 1},
    {'params': model.linear.parameters(), 'lr': 0.1},
], lr=3e-4, momentum=0.9)
````
Such as this, the optimizer has two different groups. Can anyone give an example?


Discussion:
cc @KickItLikeShika
@xiaoye-hhh thanks for question, you can check https://pytorch-ignite.ai/how-to-guides/04-fastai-lr-finder/#with-lr-finder and update it to multiple groups.

Our FastaiLRFinder can accept multiple groups but checks with a single lr range without respecting initial lr (e.g. 1.0, 0.1) in your case. So, it means that it will most probably suggest the same lr for both groups.

```python
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.models import resnet18
from torchvision.transforms import Compose, Normalize, ToTensor

from ignite.engine import create_supervised_trainer, create_supervised_evaluator
from ignite.metrics import Accuracy, Loss
from ignite.handlers import FastaiLRFinder

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()

        self.model = resnet18(num_classes=20)
        self.model.conv1 = nn.Conv2d(
            1, 64, kernel_size=3, padding=1, bias=False
        )
        self.linear = nn.Linear(20, 10)

    def forward(self, x):
        return self.model(x)


model = Net().to(device)

data_transform = Compose([ToTensor(), Normalize((0.1307,), (0.3081,))])

train_loader = DataLoader(
    MNIST(download=True, root=".", transform=data_transform, train=True),
    batch_size=128,
    shuffle=True,
)

test_loader = DataLoader(
    MNIST(download=True, root=".", transform=data_transform, train=False),
    batch_size=256,
    shuffle=False,
)


model = Net().to(device)
# optimizer = torch.optim.RMSprop(model.parameters(), lr=1e-06)
optimizer = torch.optim.SGD([
    {'params': model.model.parameters(), 'lr': 0.1},
    {'params': model.linear.parameters(), 'lr': 0.01},
], momentum=0.9)

criterion = nn.CrossEntropyLoss()

trainer = create_supervised_trainer(model, optimizer, criterion, device=device)
lr_finder = FastaiLRFinder()

# To restore the model's and optimizer's states after running the LR Finder
to_save = {"model": model, "optimizer": optimizer}

with lr_finder.attach(trainer, to_save, end_lr=1.0) as trainer_with_lr_finder:
    trainer_with_lr_finder.run(train_loader)

print("Suggested LR", lr_finder.lr_suggestion())
> Suggested LR [0.10451768106330113, 0.10451768106330113]
```
Thanks for quick reply. As your code show, the lr_suggestion of the two groups are both 0.1045. However when I use pretrained model and do some change at last, I want the lrs of different groups are different, smaller for pretrained and bigger for new. How to do is?
@xiaoye-hhh right now there is no way to perform a search for 2 params groups. We can try to expand the behaviour of our lr finder to accept list of floats for start lr and end lr. It wont be a ND search (where N is the number of groups) as it may be too expensive.

Here is how you can patch current code to make it work with 2 groups:
```python
from ignite.engine import Engine, Events
from typing import Union, List
from ignite.handlers.lr_finder import _LRScheduler
from ignite.handlers.param_scheduler import LRScheduler, PiecewiseLinear


class UpdatedLRFinder(FastaiLRFinder):

  def _run(
        self,
        trainer,
        optimizer,
        output_transform,
        num_iter: int,
        start_lr: Union[None, float, List[float]],
        end_lr: Union[float, List[float]],
        step_mode: str,
        smooth_f: float,
        diverge_th: float,
    ) -> None:

        self._history = {"lr": [], "loss": []}
        self._best_loss = None
        self._diverge_flag = False

        # attach LRScheduler to trainer.
        if num_iter is None:
            num_iter = trainer.state.epoch_length * trainer.state.max_epochs
        else:
            max_iter = trainer.state.epoch_length * trainer.state.max_epochs  # type: ignore[operator]
            if max_iter < num_iter:
                max_iter = num_iter
                trainer.state.max_epochs = ceil(num_iter / trainer.state.epoch_length)  # type: ignore[operator]

        if not trainer.has_event_handler(self._reached_num_iterations):
            trainer.add_event_handler(Events.ITERATION_COMPLETED, self._reached_num_iterations, num_iter)

        # attach loss and lr logging
        if not trainer.has_event_handler(self._log_lr_and_loss):
            trainer.add_event_handler(
                Events.ITERATION_COMPLETED, self._log_lr_and_loss, output_transform, smooth_f, diverge_th
            )

        self.logger.debug(f"Running LR finder for {num_iter} iterations")
        # Initialize the proper learning rate policy
        if step_mode.lower() == "exp":
            if start_lr is None:
                start_lr_list = [optimizer.param_groups[i]["lr"] for i in range(len(optimizer.param_groups))]
            elif isinstance(start_lr, float):
                start_lr_list = [start_lr] * len(optimizer.param_groups)
            else:
                assert len(start_lr) == len(optimizer.param_groups)
                start_lr_list = start_lr

            if isinstance(end_lr, float):
                end_lr_list = [end_lr] * len(optimizer.param_groups)
            else:
                assert len(end_lr) == len(optimizer.param_groups)
                end_lr_list = end_lr
            
            self._lr_schedule = LRScheduler(_ExponentialLR(optimizer, start_lr_list, end_lr_list, num_iter))
        else:
            if isinstance(start_lr, list) or isinstance(end_lr, list):
                assert False, "THIS WONT WORK"
            self._lr_schedule = PiecewiseLinear(
                optimizer, param_name="lr", milestones_values=[(0, start_lr), (num_iter, end_lr)]
            )
        if not trainer.has_event_handler(self._lr_schedule):
            trainer.add_event_handler(Events.ITERATION_COMPLETED, self._lr_schedule, num_iter)


class _ExponentialLR(_LRScheduler): 
    def __init__(self, optimizer, start_lrs: List[float], end_lrs: List[float], num_iter: int, last_epoch: int = -1):
        self.end_lrs = end_lrs
        self.num_iter = num_iter
        super(_ExponentialLR, self).__init__(optimizer, last_epoch)

        # override base_lrs
        self.base_lrs = start_lrs

    def get_lr(self) -> List[float]:  # type: ignore
        curr_iter = self.last_epoch + 1
        r = curr_iter / self.num_iter
        return [base_lr * (end_lr / base_lr) ** r for end_lr, base_lr in zip(self.end_lrs, self.base_lrs)]


lr_finder = UpdatedLRFinder()
with lr_finder.attach(trainer, to_save, start_lr=[0.001, 0.01], end_lr=[0.1, 2.0]) as trainer_with_lr_finder:
    trainer_with_lr_finder.run(train_loader)

print("Suggested LR", lr_finder.lr_suggestion())
> Suggested LR [0.0011473647220774028, 0.01171352105579639]
```
![image](https://user-images.githubusercontent.com/2459423/189089042-b43a7987-a848-4455-93d4-6023426d9809.png)



Maybe, @Jacob208M will help to integrate a similar feature to the library.

Thanks a lot. It's helpful.

Resolving patch:
diff --git a/ignite/handlers/lr_finder.py b/ignite/handlers/lr_finder.py
--- a/ignite/handlers/lr_finder.py
+++ b/ignite/handlers/lr_finder.py
@@ -14,7 +14,7 @@
 import ignite.distributed as idist
 from ignite.engine import Engine, Events
 from ignite.handlers import Checkpoint
-from ignite.handlers.param_scheduler import LRScheduler, PiecewiseLinear
+from ignite.handlers.param_scheduler import LRScheduler, ParamGroupScheduler, PiecewiseLinear
 
 
 class FastaiLRFinder:
@@ -74,11 +74,12 @@ class FastaiLRFinder:
     .. versionadded:: 0.4.6
     """
 
+    _lr_schedule: Union[LRScheduler, PiecewiseLinear, ParamGroupScheduler]
+
     def __init__(self) -> None:
         self._diverge_flag = False
         self._history = {}  # type: Dict[str, List[Any]]
         self._best_loss = None
-        self._lr_schedule = None  # type: Optional[Union[LRScheduler, PiecewiseLinear]]
         self.logger = logging.getLogger(__name__ + "." + self.__class__.__name__)
 
     def _run(
@@ -87,8 +88,8 @@ def _run(
         optimizer: Optimizer,
         output_transform: Callable,
         num_iter: int,
-        start_lr: float,
-        end_lr: float,
+        start_lrs: List[float],
+        end_lrs: List[float],
         step_mode: str,
         smooth_f: float,
         diverge_th: float,
@@ -118,22 +119,35 @@ def _run(
             )
 
         self.logger.debug(f"Running LR finder for {num_iter} iterations")
-        if start_lr is None:
-            start_lr = optimizer.param_groups[0]["lr"]
+
         # Initialize the proper learning rate policy
         if step_mode.lower() == "exp":
-            start_lr_list = [start_lr] * len(optimizer.param_groups)
-            self._lr_schedule = LRScheduler(_ExponentialLR(optimizer, start_lr_list, end_lr, num_iter))
+            self._lr_schedule = LRScheduler(_ExponentialLR(optimizer, start_lrs, end_lrs, num_iter))
         else:
-            self._lr_schedule = PiecewiseLinear(
-                optimizer, param_name="lr", milestones_values=[(0, start_lr), (num_iter, end_lr)]
-            )
+            if len(start_lrs) == 1:
+                self._lr_schedule = PiecewiseLinear(
+                    optimizer,
+                    param_name="lr",
+                    milestones_values=[(0, start_lrs[0]), (num_iter, end_lrs[0])],
+                )
+            else:
+                self._lr_schedule = ParamGroupScheduler(
+                    [
+                        PiecewiseLinear(
+                            optimizer,
+                            param_name="lr",
+                            milestones_values=[(0, start_lrs[i]), (num_iter, end_lrs[i])],
+                            param_group_index=i,
+                        )
+                        for i in range(len(optimizer.param_groups))
+                    ]
+                )
         if not trainer.has_event_handler(self._lr_schedule):
             trainer.add_event_handler(Events.ITERATION_COMPLETED, self._lr_schedule, num_iter)
 
     def _reset(self, trainer: Engine) -> None:
         self.logger.debug("Completed LR finder run")
-        trainer.remove_event_handler(self._lr_schedule, Events.ITERATION_COMPLETED)  # type: ignore[arg-type]
+        trainer.remove_event_handler(self._lr_schedule, Events.ITERATION_COMPLETED)
         trainer.remove_event_handler(self._log_lr_and_loss, Events.ITERATION_COMPLETED)
         trainer.remove_event_handler(self._reached_num_iterations, Events.ITERATION_COMPLETED)
 
@@ -157,7 +171,7 @@ def _log_lr_and_loss(self, trainer: Engine, output_transform: Callable, smooth_f
                     f"but got output of type {type(loss).__name__}"
                 )
         loss = idist.all_reduce(loss)
-        lr = self._lr_schedule.get_param()  # type: ignore[union-attr]
+        lr = self._lr_schedule.get_param()
         self._history["lr"].append(lr)
         if trainer.state.iteration == 1:
             self._best_loss = loss
@@ -251,7 +265,6 @@ def plot(
             )
         if not self._history:
             raise RuntimeError("learning rate finder didn't run yet so results can't be plotted")
-
         if skip_start < 0:
             raise ValueError("skip_start cannot be negative")
         if skip_end < 0:
@@ -367,8 +380,8 @@ def attach(
         to_save: Mapping,
         output_transform: Callable = lambda output: output,
         num_iter: Optional[int] = None,
-        start_lr: Optional[float] = None,
-        end_lr: float = 10.0,
+        start_lr: Optional[Union[float, List[float]]] = None,
+        end_lr: Optional[Union[float, List[float]]] = 10.0,
         step_mode: str = "exp",
         smooth_f: float = 0.05,
         diverge_th: float = 5.0,
@@ -432,8 +445,37 @@ def attach(
                 raise TypeError(f"if provided, num_iter should be an integer, but give {num_iter}")
             if num_iter <= 0:
                 raise ValueError(f"if provided, num_iter should be positive, but give {num_iter}")
-        if isinstance(start_lr, (float, int)) and start_lr >= end_lr:
-            raise ValueError(f"start_lr must be less than end_lr, start_lr={start_lr} vs end_lr={end_lr}")
+
+        optimizer = to_save["optimizer"]
+        if start_lr is None:
+            start_lrs = [pg["lr"] for pg in optimizer.param_groups]
+        elif isinstance(start_lr, float):
+            start_lrs = [start_lr] * len(optimizer.param_groups)
+        elif isinstance(start_lr, list):
+            if len(start_lr) != len(optimizer.param_groups):
+                raise ValueError(
+                    "Number of values of start_lr should be equal to optimizer values."
+                    f"start_lr values:{len(start_lr)} optimizer values: {len(optimizer.param_groups)}"
+                )
+            start_lrs = start_lr
+        else:
+            raise TypeError(f"start_lr should be a float or list of floats, but given {type(start_lr)}")
+
+        if isinstance(end_lr, float):
+            end_lrs = [end_lr] * len(optimizer.param_groups)
+        elif isinstance(end_lr, list):
+            if len(end_lr) != len(optimizer.param_groups):
+                raise ValueError(
+                    "Number of values of end_lr should be equal to optimizer values."
+                    f"end_lr values:{len(end_lr)} optimizer values: {len(optimizer.param_groups)}"
+                )
+            end_lrs = end_lr
+        else:
+            raise TypeError(f"end_lr should be a float or list of floats, but given {type(end_lr)}")
+
+        for start, end in zip(start_lrs, end_lrs):
+            if start >= end:
+                raise ValueError(f"start_lr must be less than end_lr, start_lr={start_lr} vs end_lr={end_lr}")
 
         # store to_save
         with tempfile.TemporaryDirectory() as tmpdirname:
@@ -443,7 +485,6 @@ def attach(
             cache_filepath = Path(tmpdirname) / "ignite_lr_finder_cache.pt"
             torch.save(obj, cache_filepath.as_posix())
 
-            optimizer = to_save["optimizer"]
             # Attach handlers
             if not trainer.has_event_handler(self._run):
                 trainer.add_event_handler(
@@ -452,8 +493,8 @@ def attach(
                     optimizer,
                     output_transform,
                     num_iter,
-                    start_lr,
-                    end_lr,
+                    start_lrs,
+                    end_lrs,
                     step_mode,
                     smooth_f,
                     diverge_th,
@@ -479,23 +520,24 @@ class _ExponentialLR(_LRScheduler):
 
     Args:
         optimizer: wrapped optimizer.
-        end_lr: the initial learning rate which is the lower
-            boundary of the test. Default: 10.
+        start_lrs: the initial learning rate for parameter groups.
+        end_lrs: the final learning rate for parameter groups.
         num_iter: the number of iterations over which the test
             occurs. Default: 100.
         last_epoch: the index of last epoch. Default: -1.
-
     """
 
-    def __init__(self, optimizer: Optimizer, start_lr: List[float], end_lr: float, num_iter: int, last_epoch: int = -1):
-        self.end_lr = end_lr
+    def __init__(
+        self, optimizer: Optimizer, start_lrs: List[float], end_lrs: List[float], num_iter: int, last_epoch: int = -1
+    ):
+        self.end_lrs = end_lrs
         self.num_iter = num_iter
         super(_ExponentialLR, self).__init__(optimizer, last_epoch)
 
         # override base_lrs
-        self.base_lrs = start_lr
+        self.base_lrs = start_lrs
 
-    def get_lr(self) -> List[float]:  # type: ignore
+    def get_lr(self) -> List[float]:  # type: ignore[override]
         curr_iter = self.last_epoch + 1
         r = curr_iter / self.num_iter
-        return [base_lr * (self.end_lr / base_lr) ** r for base_lr in self.base_lrs]
+        return [base_lr * (end_lr / base_lr) ** r for end_lr, base_lr in zip(self.end_lrs, self.base_lrs)]


Test patch:
diff --git a/tests/ignite/handlers/test_lr_finder.py b/tests/ignite/handlers/test_lr_finder.py
--- a/tests/ignite/handlers/test_lr_finder.py
+++ b/tests/ignite/handlers/test_lr_finder.py
@@ -161,32 +161,32 @@ def mnist_dataloader():
 def test_attach_incorrect_input_args(lr_finder, dummy_engine, model, optimizer, dataloader):
 
     with pytest.raises(TypeError, match=r"Argument to_save should be a mapping"):
-        with lr_finder.attach(dummy_engine, to_save=123) as f:
+        with lr_finder.attach(dummy_engine, to_save=123):
             pass
 
     with pytest.raises(TypeError, match=r"Object <class 'int'> should have `state_dict` method"):
-        with lr_finder.attach(dummy_engine, to_save={1: 2}) as f:
+        with lr_finder.attach(dummy_engine, to_save={1: 2}):
             pass
 
     with pytest.raises(ValueError, match=r"Mapping to_save should contain 'optimizer' key"):
-        with lr_finder.attach(dummy_engine, to_save={"model": model}) as f:
+        with lr_finder.attach(dummy_engine, to_save={"model": model}):
             pass
 
     to_save = {"model": model, "optimizer": optimizer}
     with pytest.raises(ValueError, match=r"smooth_f is outside the range \[0, 1\]"):
-        with lr_finder.attach(dummy_engine, to_save=to_save, smooth_f=234) as f:
+        with lr_finder.attach(dummy_engine, to_save=to_save, smooth_f=234):
             pass
 
     with pytest.raises(ValueError, match=r"diverge_th should be larger than 1"):
-        with lr_finder.attach(dummy_engine, to_save=to_save, diverge_th=0.0) as f:
+        with lr_finder.attach(dummy_engine, to_save=to_save, diverge_th=0.0):
             pass
 
     with pytest.raises(TypeError, match=r"if provided, num_iter should be an integer"):
-        with lr_finder.attach(dummy_engine, to_save=to_save, num_iter=0.0) as f:
+        with lr_finder.attach(dummy_engine, to_save=to_save, num_iter=0.0):
             pass
 
     with pytest.raises(ValueError, match=r"if provided, num_iter should be positive"):
-        with lr_finder.attach(dummy_engine, to_save=to_save, num_iter=0) as f:
+        with lr_finder.attach(dummy_engine, to_save=to_save, num_iter=0):
             pass
 
     with pytest.raises(TypeError, match=r"Object to_save\['optimizer'] should be torch optimizer"):
@@ -194,7 +194,7 @@ def test_attach_incorrect_input_args(lr_finder, dummy_engine, model, optimizer,
             pass
 
     with pytest.raises(ValueError, match=r"step_mode should be 'exp' or 'linear'"):
-        with lr_finder.attach(dummy_engine, to_save=to_save, step_mode="abc") as f:
+        with lr_finder.attach(dummy_engine, to_save=to_save, step_mode="abc"):
             pass
 
     with lr_finder.attach(dummy_engine, to_save) as trainer_with_finder:
@@ -205,6 +205,20 @@ def test_attach_incorrect_input_args(lr_finder, dummy_engine, model, optimizer,
     with pytest.raises(ValueError, match=r"skip_end cannot be negative"):
         lr_finder.plot(skip_end=-1)
 
+    with pytest.raises(ValueError, match=r"Number of values of start_lr should be equal to optimizer values."):
+        with lr_finder.attach(dummy_engine, to_save, start_lr=[0.1, 0.1]):
+            pass
+    with pytest.raises(ValueError, match=r"Number of values of end_lr should be equal to optimizer values."):
+        with lr_finder.attach(dummy_engine, to_save, end_lr=[0.1, 0.1]):
+            pass
+
+    with pytest.raises(TypeError, match=r"start_lr should be a float or list of floats"):
+        with lr_finder.attach(dummy_engine, to_save, start_lr=1):
+            pass
+    with pytest.raises(TypeError, match=r"end_lr should be a float or list of floats"):
+        with lr_finder.attach(dummy_engine, to_save, end_lr=1):
+            pass
+
 
 def test_attach_without_with(lr_finder, dummy_engine, to_save):
     _ = lr_finder.attach(dummy_engine, to_save=to_save)
@@ -232,15 +246,22 @@ def test_with_attach(lr_finder, to_save, dummy_engine, dataloader):
         assert len(dummy_engine._event_handlers[event]) == 0
 
 
-def test_wrong_values_start_lr_and_end_lr(lr_finder, to_save, dummy_engine, dataloader):
+def test_wrong_values_start_lr_and_end_lr(
+    lr_finder, dummy_engine, to_save, dummy_engine_mulitple_param_groups, to_save_mulitple_param_groups
+):
 
     with pytest.raises(ValueError, match=r"start_lr must be less than end_lr"):
-        with lr_finder.attach(dummy_engine, to_save=to_save, start_lr=10, end_lr=1) as trainer_with_finder:
-            trainer_with_finder.run(dataloader)
+        with lr_finder.attach(dummy_engine, to_save=to_save, start_lr=10.0, end_lr=1.0):
+            pass
 
     with pytest.raises(ValueError, match=r"start_lr must be less than end_lr"):
-        with lr_finder.attach(dummy_engine, to_save=to_save, start_lr=10, end_lr=10) as trainer_with_finder:
-            trainer_with_finder.run(dataloader)
+        with lr_finder.attach(
+            dummy_engine_mulitple_param_groups,
+            to_save=to_save_mulitple_param_groups,
+            start_lr=[1.0, 10.0, 5.0],
+            end_lr=[10.0, 10.0, 10.0],
+        ):
+            pass
 
 
 def test_model_optimizer_reset(lr_finder, to_save, dummy_engine, dataloader):
@@ -275,6 +296,24 @@ def test_lr_policy(lr_finder, to_save, dummy_engine, dataloader):
     assert all([lr[i - 1] < lr[i] for i in range(1, len(lr))])
 
 
+@pytest.mark.parametrize("step_mode", ["exp", "linear"])
+def test_multiple_optimizers(
+    lr_finder, dummy_engine_mulitple_param_groups, to_save_mulitple_param_groups, dataloader, step_mode
+):
+    start_lr = [0.1, 0.1, 0.01]
+    end_lr = [1.0, 1.0, 1.0]
+    with lr_finder.attach(
+        dummy_engine_mulitple_param_groups,
+        to_save_mulitple_param_groups,
+        start_lr=start_lr,
+        end_lr=end_lr,
+        step_mode=step_mode,
+    ) as trainer:
+        trainer.run(dataloader)
+    groups_lrs = lr_finder.get_results()["lr"]
+    assert [all([group_lrs[i - 1] < group_lrs[i] for i in range(1, len(group_lrs))]) for group_lrs in groups_lrs]
+
+
 def assert_output_sizes(lr_finder, dummy_engine):
     iteration = dummy_engine.state.iteration
     lr_finder_results = lr_finder.get_results()
@@ -313,7 +352,7 @@ def test_num_iter_is_not_enough(lr_finder, to_save, dummy_engine, dataloader):
 
 
 def test_detach_terminates(lr_finder, to_save, dummy_engine, dataloader):
-    with lr_finder.attach(dummy_engine, to_save, end_lr=100, diverge_th=2) as trainer_with_finder:
+    with lr_finder.attach(dummy_engine, to_save, end_lr=100.0, diverge_th=2) as trainer_with_finder:
         trainer_with_finder.run(dataloader)
 
     dummy_engine.run(dataloader, max_epochs=3)
@@ -335,7 +374,7 @@ def test_different_num_iters(lr_finder, to_save, dummy_engine, dataloader):
 @pytest.mark.parametrize("step_mode", ["exp", "linear"])
 def test_start_lr(lr_finder, to_save, dummy_engine, dataloader, step_mode):
     with lr_finder.attach(
-        dummy_engine, to_save, start_lr=0.01, end_lr=10, num_iter=5, step_mode=step_mode, diverge_th=1
+        dummy_engine, to_save, start_lr=0.01, end_lr=10.0, num_iter=5, step_mode=step_mode, diverge_th=1
     ) as trainer_with_finder:
         trainer_with_finder.run(dataloader)
     history = lr_finder.get_results()
@@ -486,7 +525,7 @@ def test_no_matplotlib(no_site_packages, lr_finder):
 
 def test_plot_single_param_group(dirname, lr_finder, mnist_to_save, dummy_engine_mnist, mnist_dataloader):
 
-    with lr_finder.attach(dummy_engine_mnist, mnist_to_save, end_lr=20, smooth_f=0.04) as trainer_with_finder:
+    with lr_finder.attach(dummy_engine_mnist, mnist_to_save, end_lr=20.0, smooth_f=0.04) as trainer_with_finder:
         trainer_with_finder.run(mnist_dataloader)
 
     def _test(ax):
@@ -516,7 +555,7 @@ def test_plot_multiple_param_groups(
 ):
 
     with lr_finder.attach(
-        dummy_engine_mulitple_param_groups, to_save_mulitple_param_groups, end_lr=20, smooth_f=0.04
+        dummy_engine_mulitple_param_groups, to_save_mulitple_param_groups, end_lr=20.0, smooth_f=0.04
     ) as trainer_with_finder:
         trainer_with_finder.run(dataloader_plot)