code: def save_checkpoint(
      iterations: int,
      epoch: int,
      model: torch.nn.Module,
      optimizer: Union[BaseOptim, torch.optim.Optimizer],
      best_metric: float,
      is_best: bool,
      save_dir: str,
      gradient_scaler: torch.cuda.amp.GradScaler,
      model_ema: Optional[torch.nn.Module] = None,
      is_ema_best: bool = False,
      ema_best_metric: Optional[float] = None,
      max_ckpt_metric: bool = False,
      k_best_checkpoints: int = -1,
      save_all_checkpoints: bool = False,
      *args,
      **kwargs,
  ) -> None:
      """Save checkpoints corresponding to the current state of the training.
  
      Args:
          iterations: An integer denoting training iteration number. Each iteration corresponds to forward-backward passes
              on a batch with all GPUs.
          epoch: An integer denoting epoch number.
          model: The model being trained.
          optimizer: Optimizer object, which possibly store training optimization state variables.
          best_metric: Best observed value of the tracking validation metric. For example, best top-1 validation accuracy
              that is observed until the current iteration.
          is_best: A boolean demonstrating whether the current model obtains the best validation metric compared
              to the previously saved checkpoints.
          save_dir: Path to a directory to save checkpoints.
          gradient_scaler: `GradScaler` object storing required automatic mixed precision state.
          model_ema: EMA model to be stored in the checkpoint.
          is_ema_best: A boolean demonstrating whether the current EMA model obtains the best validation metric compared
              to the previously saved checkpoints.
          ema_best_metric: Best observed value of the tracking validation metric by the EMA model.
          max_ckpt_metric: A boolean demonstrating whether the tracking validation metric is higher the better, or lower
              the better.
          k_best_checkpoints: An integer k determining number of top (based on validation metric) checkpoints to keep.
              If `k_best_checkpoints` is smaller than 1, only best checkpoint is stored.
          save_all_checkpoints: If True, will save model_state checkpoints (main model and its EMA) for all epochs.
      """
      checkpoint = get_training_state(
          iterations,
          epoch,
          model,
          optimizer,
          best_metric,
          gradient_scaler,
          model_ema,
      )
      model_state = checkpoint.get("model_state_dict")
      ckpt_str = "{}/checkpoint".format(save_dir)
  
      if is_best:
          best_model_fname = "{}_best.{}".format(
              ckpt_str, CHECKPOINT_EXTN
          )
          if os.path.isfile(best_model_fname):
              os.remove(best_model_fname)
          torch.save(model_state, best_model_fname)
          logger.log(
              "Best checkpoint with score {:.2f} saved at {}".format(
                  best_metric, best_model_fname
              )
          )
  
          if k_best_checkpoints > 1:
              avg_and_save_k_checkpoints(
                  model_state,
                  best_metric,
                  k_best_checkpoints,
                  max_ckpt_metric,
                  ckpt_str,
              )
  
      if model_ema is not None:
          ema_fname = "{}_ema_last.{}".format(ckpt_str, CHECKPOINT_EXTN)
          torch.save(checkpoint["ema_state_dict"], ema_fname)
          if save_all_checkpoints:
              ema_fname = "{}_ema_epoch{}.{}".format(
                  ckpt_str, epoch, CHECKPOINT_EXTN
              )
              torch.save(checkpoint["ema_state_dict"], ema_fname)
  
          if is_ema_best:
              ema_best_fname = "{}_ema_best.{}".format(
                  ckpt_str, CHECKPOINT_EXTN
              )
              if os.path.isfile(ema_best_fname):
                  os.remove(ema_best_fname)
              torch.save(checkpoint["ema_state_dict"], ema_best_fname)
              logger.log(
                  "Best EMA checkpoint with score {:.2f} saved at {}".format(
                      ema_best_metric, ema_best_fname
                  )
              )
  
              if k_best_checkpoints > 1 and ema_best_metric is not None:
                  avg_and_save_k_checkpoints(
                      model_state=checkpoint["ema_state_dict"],
                      best_metric=ema_best_metric,
                      k_best_checkpoints=k_best_checkpoints,
                      max_ckpt_metric=max_ckpt_metric,
                      ckpt_str="{}_ema".format(ckpt_str),
                  )
  
      ckpt_fname = "{}/training_checkpoint_last.{}".format(
          save_dir, CHECKPOINT_EXTN
      )
      torch.save(checkpoint, ckpt_fname)
  
      ckpt_fname = "{}_last.{}".format(ckpt_str, CHECKPOINT_EXTN)
      torch.save(model_state, ckpt_fname)
  
      if save_all_checkpoints:
          ckpt_fname = "{}_epoch{}.{}".format(
              ckpt_str, epoch, CHECKPOINT_EXTN
          )
          torch.save(model_state, ckpt_fname)
  
imports: import argparse
  import glob
  import math
  import os
  from typing import Dict, List, Optional, Tuple, Union
  import torch
  from cvnets import EMA
  from optim import BaseOptim
  from utils import logger
  from utils.common_utils import unwrap_model_fn
  from utils.ddp_utils import is_master
  from utils.download_utils import get_local_path
label_desc: Use image node_id 7 for calling a global function with example usage: save_checkpoint(iterations, epoch, model, optimizer, best_metric, is_best, save_dir, gradient_scaler, model_ema, is_ema_best, ema_best_metric, max_ckpt_metric, k_best_checkpoints, save_all_checkpoints) without return types