code: def main(cfg):
      name_prefix, checkpoint_paths, save_ckpt_only = process_config(
          cfg
      )
  
      if not save_ckpt_only:
          trainer = pl.Trainer(**cfg.trainer)
  
          # <<< Change model class here ! >>>
          # Model architecture which will contain the averaged checkpoints
          # Change the model constructor to the one you would like (if needed)
          model = EncDecCTCModelBPE(cfg=cfg.model, trainer=trainer)
  
      """ < Checkpoint Averaging Logic > """
      # load state dicts
      n = len(checkpoint_paths)
      avg_state = None
  
      logging.info(f"Averaging {n} checkpoints ...")
  
      for ix, path in enumerate(checkpoint_paths):
          checkpoint = torch.load(path, map_location="cpu")
  
          if "state_dict" in checkpoint:
              checkpoint = checkpoint["state_dict"]
  
          if ix == 0:
              # Initial state
              avg_state = checkpoint
  
              logging.info(
                  f"Initialized average state dict with checkpoint : {path}"
              )
          else:
              # Accumulated state
              for k in avg_state:
                  avg_state[k] = avg_state[k] + checkpoint[k]
  
              logging.info(
                  f"Updated average state dict with state from checkpoint : {path}"
              )
  
      for k in avg_state:
          if str(avg_state[k].dtype).startswith("torch.int"):
              # For int type, not averaged, but only accumulated.
              # e.g. BatchNorm.num_batches_tracked
              pass
          else:
              avg_state[k] = avg_state[k] / n
  
      # Save model
      if save_ckpt_only:
          ckpt_name = name_prefix + "-averaged.ckpt"
          torch.save(avg_state, ckpt_name)
  
          logging.info(
              f"Averaged pytorch checkpoint saved as : {ckpt_name}"
          )
      else:
          # Set model state
          logging.info("Loading averaged state dict in provided model")
          model.load_state_dict(avg_state, strict=True)
  
          ckpt_name = name_prefix + "-averaged.nemo"
          model.save_to(ckpt_name)
  
          logging.info(f"Averaged model saved as : {ckpt_name}")
  
file_path: NeMo/scripts/checkpoint_averaging/average_model_checkpoints.py
imports: import os
  import pytorch_lightning
  import torch
  from omegaconf import OmegaConf, open_dict
  from nemo.collections.asr.models import EncDecCTCModelBPE
  from nemo.core.config import hydra_runner
  from nemo.utils import logging
label_desc: Use image node_id 2 for calling a global function with example usage: main(cfg) without return types