docstr: Test to ensure that when we re-load the EMA callback, it loads the EMA weights correctly.
code: def test_ema_saved_state(self, tmpdir, caplog):
      """Test to ensure that when we re-load the EMA callback, it loads the EMA weights correctly."""
      temp_path = os.path.join(tmpdir, "saved_state")
  
      class TerminateCallback(Callback):
          def on_train_epoch_end(
              self,
              trainer: "pl.Trainer",
              pl_module: "pl.LightningModule",
          ) -> None:
              self.saved_ema_weights = extract_ema_weights(
                  pl_module, trainer
              )
              self.pl_module_weights = extract_weights(pl_module)
              raise SystemExit
  
      model = ExampleModel()
      terminate_callback = TerminateCallback()
  
      trainer = Trainer(
          max_epochs=2,
          limit_val_batches=1,
          limit_train_batches=16,
          logger=False,
          val_check_interval=0.5,
          enable_checkpointing=False,
          accelerator="gpu",
          devices=1,
          callbacks=[terminate_callback],
      )
      exp_manager(
          trainer,
          {
              "ema": {"enable": True},
              "explicit_log_dir": str(temp_path),
              "checkpoint_callback_params": {
                  "filename": f"{{epoch}}-{{step}}"
              },
          },
      )
      with pytest.raises(SystemExit):
          trainer.fit(model=model)
      resume_path = os.path.join(
          temp_path, "checkpoints/epoch=0-step=8.ckpt"
      )
  
      model = ExampleModel()
  
      class CheckStateCallback(Callback):
          def on_train_start(
              self,
              trainer: "pl.Trainer",
              pl_module: "pl.LightningModule",
          ) -> None:
              weights = extract_weights(pl_module)
              for x, y in zip(
                  weights, terminate_callback.pl_module_weights
              ):
                  assert torch.allclose(x.cpu(), y.cpu())
              current_ema_weights = extract_ema_weights(
                  pl_module, trainer
              )
              for x, y in zip(
                  current_ema_weights,
                  terminate_callback.saved_ema_weights,
              ):
                  assert torch.allclose(x.cpu(), y.cpu())
  
              for optimizer in trainer.optimizers:
                  assert isinstance(optimizer, EMAOptimizer)
                  assert optimizer.current_step == 8
  
      trainer = Trainer(
          max_epochs=2,
          limit_val_batches=0,
          limit_train_batches=16,
          logger=False,
          enable_checkpointing=False,
          accelerator="gpu",
          devices=1,
      )
      exp_manager(
          trainer,
          {
              "ema": {"enable": True},
              "explicit_log_dir": str(temp_path),
              "checkpoint_callback_params": {
                  "filename": f"{{epoch}}-{{step}}"
              },
          },
      )
      # add the callback after the exp manager has made modifications.
      trainer.callbacks.append(CheckStateCallback())
      trainer.fit(model, ckpt_path=resume_path)
  
      # ensure we can resume from the EMA weights
      ema_path = os.path.join(
          temp_path, "checkpoints/epoch=0-step=8-EMA.ckpt"
      )
  
      trainer = Trainer(
          max_epochs=1,
          limit_val_batches=0,
          limit_train_batches=1,
          logger=False,
          enable_checkpointing=False,
          accelerator="gpu",
          devices=1,
      )
      exp_manager(
          trainer,
          {
              "ema": {"enable": True},
              "explicit_log_dir": str(temp_path),
              "checkpoint_callback_params": {
                  "filename": f"{{epoch}}-{{step}}"
              },
          },
      )
      trainer.fit(model, ckpt_path=ema_path)
  
      # ensure that we warn when the EMA weights do not exist
      os.remove(ema_path)
  
      trainer = Trainer(
          max_epochs=1,
          limit_val_batches=0,
          limit_train_batches=1,
          logger=False,
          enable_checkpointing=False,
          accelerator="gpu",
          devices=1,
      )
      exp_manager(
          trainer,
          {
              "ema": {
                  "enable": True,
                  "validate_original_weights": True,
              },
              "explicit_log_dir": str(temp_path),
              "checkpoint_callback_params": {
                  "filename": f"{{epoch}}-{{step}}"
              },
          },
      )
      with pytest.raises(
          MisconfigurationException,
          match="Unable to find the associated EMA weights when re-loading",
      ):
          trainer.fit(model, ckpt_path=resume_path)
  
imports: import os.path
  from typing import Any, Dict, Union
  import pytest
  import pytorch_lightning
  import torch
  from omegaconf import DictConfig, OmegaConf
  from pytorch_lightning import Callback, Trainer
  from pytorch_lightning.utilities.exceptions import MisconfigurationException
  from pytorch_lightning.utilities.types import STEP_OUTPUT
  from nemo.collections.common.callbacks import EMA
  from nemo.collections.common.callbacks.ema import EMAOptimizer
  from nemo.core import ModelPT
  from nemo.utils.exp_manager import exp_manager
label_desc: Use image node_id 2 for calling the TestEMAConfig obj's underlying member method code with example usage: obj.test_ema_saved_state(tmpdir, caplog) without return types