code: def save_checkpoint(
      file,
      epoch: int,
      model: torch.nn.Module,
      optimizer: torch.optim.Optimizer = None,
      lr_scheduler: torch.optim.lr_scheduler._LRScheduler = None,
      **kwargs,
  ):
      """Stores the checkpoint to disk. Saves all the training components' parameters or buffers, such as model, optimizer,
      lr_scheduler etc. into a checkpoint dictionary.
  
      Args:
          file: a file-like object (has to implement write and flush) or a string or os.PathLike object containing a
              file name.
          epoch (int): Epoch number (indicates how many epochs have you trained this model).
          model (:class:`torch.nn.Module`): Model to be saved.
          optimizer (Union[:class:`torch.optim.Optimizer`, :class:`colossalai.nn.optimizer`]): Optimizer to be saved.
          lr_scheduler (Union[:class:`torch.optim.lr_scheduler`, :class:`colossalai.nn.lr_scheduler`], optional):
              lr_scheduler to be saved, defaults to None.
          pickle_module: module used for pickling metadata and objects
          pickle_protocol: can be specified to override the default protocol
      """
      # ckpt container
      checkpoint = {"epoch": epoch}
  
      model_state = model.state_dict()
      if is_using_pp() and gpc.get_local_rank(ParallelMode.TENSOR) == 0:
          model_state = gather_pipeline_parallel_state_dict(model_state)
  
      if gpc.get_global_rank() == 0:
          checkpoint["model"] = model_state
  
          # if optimizer is not None:
          #     checkpoint['optimizer'] = optimizer.state_dict()
  
          # if lr_scheduler is not None:
          #     checkpoint['lr_scheduler'] = lr_scheduler.state_dict()
  
          torch.save(checkpoint, file, **kwargs)
  
file_path: ColossalAI/colossalai/legacy/utils/checkpointing.py
imports: from collections import OrderedDict
  from itertools import chain
  import torch
  import torch.distributed
  from colossalai.legacy.constants import IS_TENSOR_PARALLEL
  from colossalai.legacy.context.parallel_mode import ParallelMode
  from colossalai.legacy.core import global_context
  from .common import is_using_pp
label_desc: Use image node_id 8 for calling a global function with example usage: save_checkpoint(file, epoch, model, optimizer, lr_scheduler) without return types