code: def save_checkpoint(model: torch.nn.Module, path: str, cfg: Any):
      """Saves a model checkpoint if the path is provided.
  
      Args:
          model: model to save
          path: path to save the checkpoint to
  
      Returns:
          Dictionary with all the keys to save
      """
  
      if cfg.environment.use_deepspeed:
          if path is not None:
              # gather model params from all ranks when using Deepspeed
              status = model.save_16bit_model(path, "checkpoint.pth")  # type: ignore
              if status:
                  if cfg.environment._local_rank == 0:
                      checkpoint = {
                          "model": torch.load(
                              os.path.join(path, "checkpoint.pth"),
                              map_location="cpu",
                          )
                      }
              else:
                  logger.warning(
                      "deepspeed.save_16bit_model didn't save the model, since"
                      " stage3_gather_16bit_weights_on_model_save=False."
                      " Saving the full checkpoint instead"
                  )
                  model.save_checkpoint(  # type: ignore
                      os.path.join(path, "ds_checkpoint")
                  )
                  if cfg.environment._local_rank == 0:
                      # load to cpu
                      state_dict = (
                          get_fp32_state_dict_from_zero_checkpoint(
                              os.path.join(path, "ds_checkpoint")
                          )
                      )
                      # save as normal checkpoint that can be loaded by `load_state_dict`
                      checkpoint = {"model": state_dict}
                      torch.save(
                          checkpoint,
                          os.path.join(path, "checkpoint.pth"),
                      )
                      shutil.rmtree(os.path.join(path, "ds_checkpoint"))
  
      else:
          if cfg.environment._local_rank == 0:
              model = unwrap_model(model)
              checkpoint = {"model": model.state_dict()}
              if path is not None:
                  torch.save(
                      checkpoint, os.path.join(path, "checkpoint.pth")
                  )
  
      if (
          cfg.environment._local_rank == 0
          and "classification_head.weight" in checkpoint["model"]
      ):
          torch.save(
              checkpoint["model"]["classification_head.weight"],
              os.path.join(path, "classification_head.pth"),
          )
  
imports: import gc
  import logging
  import os
  import re
  import shutil
  from collections import OrderedDict
  from typing import Any, Dict
  import coolname
  import deepspeed
  import numpy
  import torch
  from deepspeed.runtime.dataloader import DeepSpeedDataLoader
  from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
  from peft import LoraConfig, PeftModel, get_peft_model
  from torch.cuda.amp import autocast
  from torch.nn.parallel import DistributedDataParallel
  from tqdm import tqdm
  from transformers import AutoConfig, AutoModel, BitsAndBytesConfig, GenerationMixin, StoppingCriteria, StoppingCriteriaList
  from transformers.pytorch_utils import Conv1D
  from transformers.utils import logging
  from llm_studio.src.datasets.text_utils import get_tokenizer
  from llm_studio.src.optimizers import Optimizers
  from llm_studio.src.schedulers import Schedulers
  from llm_studio.src.utils.config_utils import NON_GENERATION_PROBLEM_TYPES
  from llm_studio.src.utils.data_utils import OrderedDistributedSampler, batch_padding, cat_batches, get_inference_batch_size
  from llm_studio.src.utils.exceptions import LLMDataException, LLMModelException
  from llm_studio.src.utils.logging_utils import TqdmToLogger
  from llm_studio.src.utils.utils import save_pickle
label_desc: Use image node_id 3 for calling a global function with example usage: save_checkpoint(model, path, cfg) without return types