docstr: Adapted from https://github.com/pytorch/pytorch/blob/v1.7.1/torch/optim/swa_utils.py#L140-L154.
docstr_tok: ["Adapted","from","https",":","\/\/github.com\/pytorch\/pytorch\/blob\/v1.7.1\/torch\/optim\/swa_utils.py","#","L140-L154","."]
code: def reset_batch_norm_and_save_state(
      self, pl_module: "pl.LightningModule"
  ) -> None:
      """Adapted from https://github.com/pytorch/pytorch/blob/v1.7.1/torch/optim/swa_utils.py#L140-L154."""
      self.momenta = {}
      for module in pl_module.modules():
          if not isinstance(module, nn.modules.batchnorm._BatchNorm):
              continue
          assert module.running_mean is not None
          module.running_mean = torch.zeros_like(
              module.running_mean,
              device=pl_module.device,
              dtype=module.running_mean.dtype,
          )
          assert module.running_var is not None
          module.running_var = torch.ones_like(
              module.running_var,
              device=pl_module.device,
              dtype=module.running_var.dtype,
          )
          self.momenta[module] = module.momentum
          module.momentum = None  # type: ignore[assignment]
          assert module.num_batches_tracked is not None
          module.num_batches_tracked *= 0
  
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 11 for calling the StochasticWeightAveraging obj's underlying member method code with example usage: obj.reset_batch_norm_and_save_state(pl_module) without return types