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_load_checkpoint(
      self,
      trainer: "pl.Trainer",
      pl_module: "pl.LightningModule",
      checkpoint: Dict[str, Any],
  ) -> None:
      checkpoint_callback = trainer.checkpoint_callback
  
      # use the connector as NeMo calls the connector directly in the exp_manager when restoring.
      connector = trainer._checkpoint_connector
      # Replace connector._ckpt_path with below to avoid calling into lightning's protected API
      ckpt_path = trainer.ckpt_path
  
      if (
          ckpt_path
          and checkpoint_callback is not None
          and "NeMo" in type(checkpoint_callback).__name__
      ):
          ext = checkpoint_callback.FILE_EXTENSION
          if ckpt_path.endswith(f"-EMA{ext}"):
              rank_zero_info(
                  "loading EMA based weights. "
                  "The callback will treat the loaded EMA weights as the main weights"
                  " and create a new EMA copy when training."
              )
              return
          ema_path = ckpt_path.replace(ext, f"-EMA{ext}")
          if os.path.exists(ema_path):
              ema_state_dict = torch.load(
                  ema_path, map_location=torch.device("cpu")
              )
  
              checkpoint["optimizer_states"] = ema_state_dict[
                  "optimizer_states"
              ]
              del ema_state_dict
              rank_zero_info("EMA state has been restored.")
          else:
              raise MisconfigurationException(
                  "Unable to find the associated EMA weights when re-loading, "
                  f"training will start with new EMA weights. Expected them to be at: {ema_path}",
              )
  
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 12 for calling the EMA obj's underlying member method code with example usage: obj.on_load_checkpoint(trainer, pl_module, checkpoint) without return types