class_docstr: EMAOptimizer is a wrapper for torch.optim.Optimizer that computes
  Exponential Moving Average of parameters registered in the optimizer.
  
  EMA parameters are automatically updated after every step of the optimizer
  with the following formula:
  
      ema_weight = decay * ema_weight + (1 - decay) * training_weight
  
  To access EMA parameters, use ``swap_ema_weights()`` context manager to
  perform a temporary in-place swap of regular parameters with EMA
  parameters.
  
  Notes:
      - EMAOptimizer is not compatible with APEX AMP O2.
  
  Args:
      optimizer (torch.optim.Optimizer): optimizer to wrap
      device (torch.device): device for EMA parameters
      decay (float): decay factor
  
  Returns:
      returns an instance of torch.optim.Optimizer that computes EMA of
      parameters
  
  Example:
      model = Model().to(device)
      opt = torch.optim.Adam(model.parameters())
  
      opt = EMAOptimizer(opt, device, 0.9999)
  
      for epoch in range(epochs):
          training_loop(model, opt)
  
          regular_eval_accuracy = evaluate(model)
  
          with opt.swap_ema_weights():
              ema_eval_accuracy = evaluate(model)
docstr: A context manager to in-place swap regular parameters with EMA
  parameters.
  It swaps back to the original regular parameters on context manager
  exit.
  
  Args:
      enabled (bool): whether the swap should be performed
docstr_tok: ["A","context","manager","to","in-place","swap","regular","parameters","with","EMA","parameters",".","It","swaps","back","to","the","original","regular","parameters","on","context","manager","exit",".","Args",":","enabled","(","bool",")",":","whether","the","swap","should","be","performed"]
code: def swap_ema_weights(self, enabled: bool = True):
      r"""
      A context manager to in-place swap regular parameters with EMA
      parameters.
      It swaps back to the original regular parameters on context manager
      exit.
  
      Args:
          enabled (bool): whether the swap should be performed
      """
  
      if enabled:
          self.switch_main_parameter_weights()
      try:
          yield
      finally:
          if enabled:
              self.switch_main_parameter_weights()
  
imports: import contextlib
  import copy
  import os
  import threading
  from typing import Any, Dict, Iterable
  import pytorch_lightning
  import torch
  from pytorch_lightning import Callback
  from pytorch_lightning.utilities.exceptions import MisconfigurationException
  from pytorch_lightning.utilities.rank_zero import rank_zero_info
label_desc: Use image node_id 8 for calling the EMAOptimizer obj's underlying member method code with example usage: obj.swap_ema_weights(enabled) without return types