class_docstr: Used to pre-train the actor model to generate better prompts.
  
  Args:
      config (ConfigActor): Configuration for the actor model
  
  Attributes:
      config (ConfigActor): Configuration for the actor model
      model (ActorModel): Actor model
      loss_function (torch.nn.CrossEntropyLoss): Loss function
      optimizer (torch.optim.Adam): Optimizer
      validation_flag (bool): Flag to indicate if the validation dataset
          is provided
      train_dataset (ActorDataset): Training dataset
      train_dataloader (DataLoader): Training dataloader
      validation_dataset (ActorDataset): Validation dataset
      validation_dataloader (DataLoader): Validation dataloader
      scheduler (torch.optim.lr_scheduler): Learning rate scheduler
      training_stats (TrainingStats): Training statistics
      model_engine (ModelEngine): Model engine for deepspeed training
      accelerator (Accelerator): Accelerator for accelerate training
  
  Methods:
      train: Train the actor model
      load_checkpoint: Load a checkpoint
      save_checkpoint: Save a checkpoint
docstr: Save the current checkpoint
  
  Args:
      current_epoch (int): Current epoch
      current_step (int): Current step
      max_epochs (int): Maximum number of epochs
      max_steps (int): Maximum number of steps
docstr_tok: ["Save","the","current","checkpoint","Args",":","current_epoch","(","int",")",":","Current","epoch","current_step","(","int",")",":","Current","step","max_epochs","(","int",")",":","Maximum","number","of","epochs","max_steps","(","int",")",":","Maximum","number","of","steps"]
code: def save_checkpoint(
      self,
      current_epoch: int,
      current_step: int,
      max_epochs: int,
      max_steps: int,
  ) -> None:
      """Save the current checkpoint
  
      Args:
          current_epoch (int): Current epoch
          current_step (int): Current step
          max_epochs (int): Maximum number of epochs
          max_steps (int): Maximum number of steps
      """
  
      print(
          f"Saving checkpoint for epoch {current_epoch + 1}, "
          f"step {current_step + 1} ..."
      )
      # look for path to save the checkpoint
      model_folder, model_name, path = ModelLoader.get_model_path(
          config=self.config,
          is_checkpoint=True,
          current_epoch=current_epoch,
          current_step=current_step,
          max_epochs=max_epochs,
          max_steps=max_steps,
      )
  
      # remove the checkpoint if it already exists
      if os.path.exists(path):
          if self.config.deepspeed_enable:
              shutil.rmtree(path)
          else:
              os.remove(path)
  
      if self.config.deepspeed_enable:
          client_state = {
              "epoch": current_epoch,
              "step": current_step,
          }
          self.model_engine.save_checkpoint(
              path, client_state=client_state
          )
      else:
          # save the checkpoint
          torch.save(
              {
                  "state_dict": self.actor.model.state_dict(),
                  "optim_state_dict": self.optimizer.state_dict(),
                  "training_stats": self.training_stats,
                  "epoch": current_epoch,
                  "step": current_step,
              },
              path,
          )
  
      # remove old checkpoints
      n_checkpoints_to_keep = self.config.n_checkpoints_to_keep
      ModelLoader.delete_old_checkpoints(
          model_folder, model_name, n_checkpoints_to_keep
      )
  
imports: import json
  import yaml
  import os
  import shutil
  import deepspeed
  import torch
  from accelerate import Accelerator
  from beartype import beartype
  from beartype.typing import Tuple
  from einops import rearrange
  from peft import get_peft_model, LoraConfig, TaskType
  from torch.utils.data import DataLoader, Dataset
  from transformers import AutoModelForCausalLM, AutoTokenizer
  from chatllama.rlhf.config import ConfigActor
  from chatllama.rlhf.model_list import hf_models_causal_lm, llama_models, hf_models
  from chatllama.rlhf.model_loader import ModelLoader
  from chatllama.rlhf.utils import TrainingStats
label_desc: Use image node_id 2 for calling the ActorTrainer obj's underlying member method code with example usage: obj.save_checkpoint(current_epoch, current_step, max_epochs, max_steps) without return types