code: def load_checkpoint(
      load_dir: Union[str, os.PathLike],
      booster: Booster,
      model: torch.nn.Module,
      optimizer: Optimizer,
      lr_scheduler: _LRScheduler,
  ) -> Tuple[int, int, int]:
      """
      Load model checkpoint, optimizer, LR scheduler and intermedidate running states.
      """
  
      # Update booster params states.
      booster.load_model(
          model=model, checkpoint=os.path.join(load_dir, "modeling")
      )
      booster.load_optimizer(
          optimizer=optimizer,
          checkpoint=os.path.join(load_dir, "optimizer"),
      )
      booster.load_lr_scheduler(
          lr_scheduler=lr_scheduler,
          checkpoint=os.path.join(load_dir, "lr_scheduler"),
      )
  
      running_states = load_json(
          file_path=os.path.join(load_dir, "running_states.json")
      )
      return (
          running_states["epoch"],
          running_states["step"],
          running_states["sample_start_index"],
      )
  
file_path: ColossalAI/applications/Colossal-LLaMA-2/colossal_llama2/utils/ckpt_io.py
imports: import json
  import os
  from typing import Any, Dict, Tuple, Union
  import torch
  from torch.optim.optimizer import Optimizer
  from torch.optim.lr_scheduler import _LRScheduler
  from colossalai.booster import Booster
  from colossalai.cluster import DistCoordinator
label_desc: Use image node_id 4 for calling a global function with example usage: load_checkpoint(load_dir, booster, model, optimizer, lr_scheduler) and returns: running_states,  running_states,  running_states