code: def on_train_epoch_start(
      self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"
  ) -> None:
      if (not self._initialized) and (
          self.swa_start <= trainer.current_epoch <= self.swa_end
      ):
          self._initialized = True
  
          # move average model to request device.
          assert self._average_model is not None
          self._average_model = self._average_model.to(
              self._device or pl_module.device
          )
  
          optimizer = trainer.optimizers[0]
          if isinstance(self._swa_lrs, float):
              self._swa_lrs = [self._swa_lrs] * len(
                  optimizer.param_groups
              )
  
          for lr, group in zip(self._swa_lrs, optimizer.param_groups):
              group["initial_lr"] = lr
  
          assert trainer.max_epochs is not None
          self._swa_scheduler = cast(
              LRScheduler,
              SWALR(
                  optimizer,
                  swa_lr=self._swa_lrs,  # type: ignore[arg-type]
                  anneal_epochs=self._annealing_epochs,
                  anneal_strategy=self._annealing_strategy,
                  last_epoch=trainer.max_epochs
                  if self._annealing_strategy == "cos"
                  else -1,
              ),
          )
          if self._scheduler_state is not None:
              # Restore scheduler state from checkpoint
              self._swa_scheduler.load_state_dict(self._scheduler_state)
          elif trainer.current_epoch != self.swa_start:
              # Log a warning if we're initializing after start without any checkpoint data,
              # as behaviour will be different compared to having checkpoint data.
              rank_zero_warn(
                  "SWA is initializing after swa_start without any checkpoint data. "
                  "This may be caused by loading a checkpoint from an older version of PyTorch Lightning."
              )
  
          # We assert that there is only one optimizer on fit start
          default_scheduler_cfg = LRSchedulerConfig(self._swa_scheduler)
          assert default_scheduler_cfg.interval == "epoch"
          assert default_scheduler_cfg.frequency == 1
  
          if trainer.lr_scheduler_configs:
              scheduler_cfg = trainer.lr_scheduler_configs[0]
              if (
                  scheduler_cfg.interval != "epoch"
                  or scheduler_cfg.frequency != 1
              ):
                  rank_zero_warn(
                      f"SWA is currently only supported every epoch. Found {scheduler_cfg}"
                  )
              rank_zero_info(
                  f"Swapping scheduler `{scheduler_cfg.scheduler.__class__.__name__}`"
                  f" for `{self._swa_scheduler.__class__.__name__}`"
              )
              trainer.lr_scheduler_configs[0] = default_scheduler_cfg
          else:
              trainer.lr_scheduler_configs.append(default_scheduler_cfg)
  
          if self.n_averaged is None:
              self.n_averaged = torch.tensor(
                  self._init_n_averaged,
                  dtype=torch.long,
                  device=pl_module.device,
              )
  
      if (self.swa_start <= trainer.current_epoch <= self.swa_end) and (
          trainer.current_epoch > self._latest_update_epoch
      ):
          assert self.n_averaged is not None
          assert self._average_model is not None
          self.update_parameters(
              self._average_model,
              pl_module,
              self.n_averaged,
              self._avg_fn,
          )
          self._latest_update_epoch = trainer.current_epoch
  
      # Note: No > here in case the callback is saved with the model and training continues
      if trainer.current_epoch == self.swa_end + 1:
          # Transfer weights from average model to pl_module
          assert self._average_model is not None
          self.transfer_weights(self._average_model, pl_module)
  
          # Reset BatchNorm for update
          self.reset_batch_norm_and_save_state(pl_module)
  
          # There is no need to perform either backward or optimizer.step as we are
          # performing only one pass over the train data-loader to compute activation statistics
          # Therefore, we will virtually increase the number of training batches by 1 and skip backward.
          trainer.fit_loop.max_batches += 1
          trainer.fit_loop._skip_backward = True
          self._accumulate_grad_batches = (
              trainer.accumulate_grad_batches
          )
          assert isinstance(
              trainer.fit_loop.max_batches, int
          ), "Iterable-style datasets are not supported"
          trainer.accumulate_grad_batches = trainer.fit_loop.max_batches
  
imports: from copy import deepcopy
  from typing import Any, Callable, Dict, List, Optional, Union, cast
  import torch
  from torch import Tensor, nn
  from torch.optim.swa_utils import SWALR
  from typing_extensions import override
  import lightning.pytorch
  from lightning.fabric.utilities.types import LRScheduler
  from lightning.pytorch.callbacks.callback import Callback
  from lightning.pytorch.strategies import DeepSpeedStrategy
  from lightning.pytorch.strategies.fsdp import FSDPStrategy
  from lightning.pytorch.utilities.exceptions import MisconfigurationException
  from lightning.pytorch.utilities.rank_zero import rank_zero_info, rank_zero_warn
  from lightning.pytorch.utilities.types import LRSchedulerConfig
label_desc: Use image node_id 7 for calling the StochasticWeightAveraging obj's underlying member method code with example usage: obj.on_train_epoch_start(trainer, pl_module) without return types