docstr: This class is a custom version of AveragedModel that allows us to skip the
  automatic deepcopy step and streamline the use of EMA / SWA. The deepcopy
  optionality gives flexibility to use modules that are not
  compatible with deepcopy, like FSDP wrapped modules. Check out
  https://github.com/pytorch/pytorch/blob/main/torch/optim/swa_utils.py#L66
  to see what the model, device, and use_buffer arguments entail.
  
  Args:
      averaging_method: Whether to use EMA or SWA.
      ema_decay: The exponential decay applied to the averaged parameters. This param
          is only needed for EMA, and is ignored otherwise (for SWA).
      skip_deepcopy: If True, will skip the deepcopy step. The user must ensure
          that the module passed in is already copied in someway
      use_lit: If True, will use Lit EMA style by adjusting weight decay based on the
          number of updates. The EMA decay will start small and will approach the
          specified ema_decay as more updates occur.
docstr_tok: ["This","class","is","a","custom","version","of","AveragedModel","that","allows","us","to","skip","the","automatic","deepcopy","step","and","streamline","the","use","of","EMA","\/","SWA",".","The","deepcopy","optionality","gives","flexibility","to","use","modules","that","are","not","compatible","with","deepcopy",",","like","FSDP","wrapped","modules",".","Check","out","https",":","\/\/github.com\/pytorch\/pytorch\/blob\/main\/torch\/optim\/swa_utils.py","#","L66","to","see","what","the","model",",","device",",","and","use_buffer","arguments","entail",".","Args",":","averaging_method",":","Whether","to","use","EMA","or","SWA",".","ema_decay",":","The","exponential","decay","applied","to","the","averaged","parameters",".","This","param","is","only","needed","for","EMA",",","and","is","ignored","otherwise","(","for","SWA",")",".","skip_deepcopy",":","If","True",",","will","skip","the","deepcopy","step",".","The","user","must","ensure","that","the","module","passed","in","is","already","copied","in","someway","use_lit",":","If","True",",","will","use","Lit","EMA","style","by","adjusting","weight","decay","based","on","the","number","of","updates",".","The","EMA","decay","will","start","small","and","will","approach","the","specified","ema_decay","as","more","updates","occur","."]
code: def __init__(
      self,
      model: torch.nn.Module,
      device: Optional[torch.device] = None,
      use_buffers: bool = False,
      averaging_method: Literal["ema", "swa"] = "ema",
      ema_decay: float = 0.999,
      skip_deepcopy: bool = False,
      use_lit: bool = False,
  ) -> None:
      """
      This class is a custom version of AveragedModel that allows us to skip the
      automatic deepcopy step and streamline the use of EMA / SWA. The deepcopy
      optionality gives flexibility to use modules that are not
      compatible with deepcopy, like FSDP wrapped modules. Check out
      https://github.com/pytorch/pytorch/blob/main/torch/optim/swa_utils.py#L66
      to see what the model, device, and use_buffer arguments entail.
  
      Args:
          averaging_method: Whether to use EMA or SWA.
          ema_decay: The exponential decay applied to the averaged parameters. This param
              is only needed for EMA, and is ignored otherwise (for SWA).
          skip_deepcopy: If True, will skip the deepcopy step. The user must ensure
              that the module passed in is already copied in someway
          use_lit: If True, will use Lit EMA style by adjusting weight decay based on the
              number of updates. The EMA decay will start small and will approach the
              specified ema_decay as more updates occur.
      """
      # setup averaging method
      if averaging_method == "ema":
          if ema_decay < 0.0 or ema_decay > 1.0:
              raise ValueError(
                  f"Decay must be between 0 and 1, got {ema_decay}"
              )
  
          # TODO: torch/optim/swa_utils.pyi needs to be updated
          # pyre-ignore Undefined attribute [16]: Module `torch.optim.swa_utils` has no attribute `get_ema_multi_avg_fn`.
          multi_avg_fn = get_ema_multi_avg_fn(ema_decay)
      elif averaging_method == "swa":
          # TODO: torch/optim/swa_utils.pyi needs to be updated
          # pyre-ignore Undefined attribute [16]: Module `torch.optim.swa_utils` has no attribute `get_swa_multi_avg_fn`.
          multi_avg_fn = get_swa_multi_avg_fn()
  
          if use_lit:
              raise ValueError("LitEMA is only supported for EMA.")
      else:
          raise ValueError(
              f"Unknown averaging method: {averaging_method}. Only ema and swa are supported."
          )
  
      self._ema_decay = ema_decay
      self._use_lit = use_lit
      self._num_updates = 0
  
      if skip_deepcopy:
          # calls parent init manually, but skips deepcopy step
          torch.nn.Module.__init__(self)  # inits grandparent class
  
          self.module: torch.nn.Module = model
          self.register_buffer(
              "n_averaged",
              torch.tensor(0, dtype=torch.long, device=device),
          )
          self.avg_fn: Optional[TSWA_avg_fn] = None
          self.multi_avg_fn: Optional[TSWA_multi_avg_fn] = multi_avg_fn
          self.use_buffers: bool = use_buffers
      else:
          # use default init implementation
  
          # TODO: torch/optim/swa_utils.pyi needs to be updated
          # pyre-ignore Unexpected keyword [28]
          super().__init__(
              model,
              device=device,
              multi_avg_fn=multi_avg_fn,
              use_buffers=use_buffers,
          )
  
imports: from typing import Callable, List, Literal, Optional
  import torch
  from torch.optim.swa_utils import AveragedModel, get_ema_multi_avg_fn, get_swa_multi_avg_fn
label_desc: Use image node_id 1 to create a new AveragedModel object from inherited base classes: PyTorchAveragedModel with example: obj = AveragedModel(model, device, use_buffers, averaging_method, ema_decay, skip_deepcopy, use_lit)