code: def test_swa_multiple_lrs(tmpdir):
      swa_lrs = [0.123, 0.321]
  
      class TestModel(BoringModel):
          def __init__(self):
              super().__init__()
              self.layer1 = torch.nn.Linear(32, 32)
              self.layer2 = torch.nn.Linear(32, 2)
              self.on_train_epoch_start_called = False
  
          def forward(self, x):
              x = self.layer1(x)
              x = self.layer2(x)
              return x
  
          def configure_optimizers(self):
              params = [
                  {"params": self.layer1.parameters(), "lr": 0.1},
                  {"params": self.layer2.parameters(), "lr": 0.2},
              ]
              return torch.optim.Adam(params)
  
          def on_train_epoch_start(self):
              optimizer = trainer.optimizers[0]
              assert [pg["lr"] for pg in optimizer.param_groups] == [
                  0.1,
                  0.2,
              ]
              assert [
                  pg["initial_lr"] for pg in optimizer.param_groups
              ] == swa_lrs
              assert [
                  pg["swa_lr"] for pg in optimizer.param_groups
              ] == swa_lrs
              self.on_train_epoch_start_called = True
  
      model = TestModel()
      swa_callback = StochasticWeightAveraging(swa_lrs=swa_lrs)
      trainer = Trainer(
          default_root_dir=tmpdir,
          callbacks=swa_callback,
          fast_dev_run=1,
      )
      trainer.fit(model)
      assert model.on_train_epoch_start_called
  
imports: import logging
  import os
  from pathlib import Path
  from typing import ContextManager, Optional
  from unittest import mock
  import pytest
  import torch
  from lightning.pytorch import Trainer
  from lightning.pytorch.callbacks import StochasticWeightAveraging
  from lightning.pytorch.demos.boring_classes import BoringModel, RandomDataset, RandomIterableDataset
  from lightning.pytorch.strategies import Strategy
  from lightning.pytorch.strategies.launchers import _MultiProcessingLauncher
  from lightning.pytorch.utilities.exceptions import MisconfigurationException
  from torch import nn
  from torch.optim.lr_scheduler import LambdaLR
  from torch.optim.swa_utils import SWALR
  from torch.utils.data import DataLoader
  from tests_pytorch.helpers.runif import RunIf
label_desc: Use image node_id 12 for calling a global function with example usage: test_swa_multiple_lrs(tmpdir) and returns: x, torch