code: def save_ckpt(
      model: torch.nn.Module,
      optimizer: Optional[torch.optim.Optimizer] = None,
      scheduler: Optional[Any] = None,
      epoch: int = 0,
  ):
      r"""Saves the model checkpoint at a given epoch."""
      ckpt: Dict[str, Any] = {}
      ckpt[MODEL_STATE] = model.state_dict()
      if optimizer is not None:
          ckpt[OPTIMIZER_STATE] = optimizer.state_dict()
      if scheduler is not None:
          ckpt[SCHEDULER_STATE] = scheduler.state_dict()
  
      os.makedirs(get_ckpt_dir(), exist_ok=True)
      torch.save(ckpt, get_ckpt_path(get_ckpt_epoch(epoch)))
  
file_path: pytorch_geometric/torch_geometric/graphgym/checkpoint.py
imports: import glob
  import os
  import os.path
  from typing import Any, Dict, List, Optional, Union
  import torch
  from torch_geometric.graphgym.config import cfg
label_desc: Use image node_id 2 for calling a global function with example usage: save_ckpt(model, optimizer, scheduler, epoch) without return types