class_docstr: Implements Exponential Moving Averaging (EMA).
  
  When training a model, this callback will maintain moving averages of the trained parameters.
  When evaluating, we use the moving averages copy of the trained parameters.
  When saving, we save an additional set of parameters with the prefix `ema`.
  
  Args:
      decay: The exponential decay used when calculating the moving average. Has to be between 0-1.
      validate_original_weights: Validate the original weights, as apposed to the EMA weights.
      every_n_steps: Apply EMA every N steps.
      cpu_offload: Offload weights to CPU.
code: def on_fit_start(
      self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"
  ) -> None:
      device = (
          pl_module.device
          if not self.cpu_offload
          else torch.device("cpu")
      )
      trainer.optimizers = [
          EMAOptimizer(
              optim,
              device=device,
              decay=self.decay,
              every_n_steps=self.every_n_steps,
              current_step=trainer.global_step,
          )
          for optim in trainer.optimizers
          if not isinstance(optim, EMAOptimizer)
      ]
  
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 2 for calling the EMA obj's underlying member method code with example usage: obj.on_fit_start(trainer, pl_module) without return types