code: def average_checkpoints(
      checkpoint_paths: List[str],
  ):
      """
      Average a list of checkpoints' state_dicts.
      Reference: https://github.com/rwightman/pytorch-image-models/blob/master/avg_checkpoints.py
  
      Parameters
      ----------
      checkpoint_paths
          A list of model checkpoint paths.
  
      Returns
      -------
      The averaged state_dict.
      """
      if len(checkpoint_paths) > 1:
          avg_state_dict = {}
          avg_counts = {}
          for per_path in checkpoint_paths:
              if os.path.isdir(
                  per_path + "-dir"
              ):  # deepspeed save checkpoints into a directory
                  from lightning.pytorch.utilities.deepspeed import (
                      convert_zero_checkpoint_to_fp32_state_dict,
                  )
  
                  convert_zero_checkpoint_to_fp32_state_dict(
                      per_path + "-dir", per_path
                  )
                  shutil.rmtree(per_path + "-dir")
                  state_dict = torch.load(
                      per_path, map_location=torch.device("cpu")
                  )["state_dict"]
              else:
                  state_dict = torch.load(
                      per_path, map_location=torch.device("cpu")
                  )["state_dict"]
              for k, v in state_dict.items():
                  if k not in avg_state_dict:
                      avg_state_dict[k] = v.clone().to(
                          dtype=torch.float64
                      )
                      avg_counts[k] = 1
                  else:
                      avg_state_dict[k] += v.to(dtype=torch.float64)
                      avg_counts[k] += 1
              del state_dict
  
          for k, v in avg_state_dict.items():
              v.div_(avg_counts[k])
  
          # convert to float32.
          float32_info = torch.finfo(torch.float32)
          for k in avg_state_dict:
              avg_state_dict[k].clamp_(
                  float32_info.min, float32_info.max
              ).to(dtype=torch.float32)
      else:
          avg_state_dict = torch.load(
              checkpoint_paths[0], map_location=torch.device("cpu")
          )["state_dict"]
  
      return avg_state_dict
  
imports: import logging
  import os
  import re
  import shutil
  from typing import Any, Dict, List, Optional, Tuple, Union
  import lightning.pytorch
  import torch
  from lightning.pytorch.strategies import DeepSpeedStrategy
  from lightning.pytorch.utilities.rank_zero import rank_zero_warn
  from .cloud_io import _atomic_save
  from .cloud_io import _load
  from .cloud_io import get_filesystem
label_desc: Use image node_id 1 for calling a global function with example usage: average_checkpoints(checkpoint_paths) and returns: avg_state_dict