code: def extract_ema_weights(pl_module, trainer):
      ema_callback = [
          x for x in trainer.callbacks if isinstance(x, EMA)
      ][0]
      ema_callback.swap_model_weights(trainer)
      weights = extract_weights(pl_module)
      ema_callback.swap_model_weights(trainer)
      return weights
  
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 1 for calling a global function with example usage: extract_ema_weights(pl_module, trainer) and returns: weights