code: def on_train_end(
      self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"
  ) -> None:
      # the trainer increases the current epoch before this hook is called
      if (
          self._model_contains_batch_norm
          and trainer.current_epoch - 1 == self.swa_end + 1
      ):
          # BatchNorm epoch update. Reset state
          trainer.accumulate_grad_batches = (
              self._accumulate_grad_batches
          )
          trainer.fit_loop.max_batches -= 1
          assert trainer.fit_loop.max_epochs is not None
          trainer.fit_loop.max_epochs -= 1
          self.reset_momenta()
      elif trainer.current_epoch - 1 == self.swa_end:
          # Last SWA epoch. Transfer weights from average model to pl_module
          assert self._average_model is not None
          self.transfer_weights(self._average_model, pl_module)
  
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 9 for calling the StochasticWeightAveraging obj's underlying member method code with example usage: obj.on_train_end(trainer, pl_module) without return types