code: def on_train_batch_end(
      self,
      trainer: "pl.Trainer",
      pl_module: "pl.LightningModule",
      outputs: STEP_OUTPUT,
      batch: Any,
      batch_idx: int,
  ) -> None:
      if (batch_idx + 1) % trainer.accumulate_grad_batches != 0:
          # skip assertion as ema weights are not updated.
          return
      ema_callback = [
          x for x in trainer.callbacks if isinstance(x, EMA)
      ][0]
      decay = ema_callback.decay
      expected_ema_weights = []
  
      new_weights = extract_weights(pl_module)
  
      for ema_weight, new_weight in zip(self.ema_weights, new_weights):
          expected_ema_weight = ema_weight * decay
          expected_ema_weight += new_weight * (1 - decay)
          expected_ema_weights.append(expected_ema_weight)
      ema_weights = extract_ema_weights(pl_module, trainer)
      for actual_ema_weight, expected_ema_weight in zip(
          ema_weights, expected_ema_weights
      ):
          assert torch.allclose(actual_ema_weight, expected_ema_weight)
      self.ema_weights = expected_ema_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 2 for calling the EMAAssertCallback obj's underlying member method code with example usage: obj.on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx) without return types