class_docstr: Class to train the reward model
  
  Args:
      config (ConfigModel): Config parameters for the model
  
  Attributes:
      model (RewardModel): Reward model
      config (ConfigModel): Config parameters for the model
      optimizer (torch.optim): Optimizer for the model
      loss_function (torch.nn): Loss function for the model
      validation_flag (bool): Flag to indicate if the validation dataset
          is available
      train_dataset (RewardDataset): Dataset for training
      validation_dataset (RewardDataset): Dataset for validation
      train_dataloader (DataLoader): Dataloader for training
      validation_dataloader (DataLoader): Dataloader for validation
      scheduler (torch.optim.lr_scheduler): Scheduler for the optimizer
      training_stats (List[Dict]): List of dictionaries with the training
          statistics
      model_engine (ModelEngine): Model engine to train the model
          using deepspeed
      accelerator (Accelerator): Accelerator to train the model using
          accelerate by HF.
  
  
  Methods:
      train: Train the reward model
      save_checkpoints: Save the checkpoints of the model
      load_checkpoints: Load the checkpoints of the model
class_docstr_tok: ["Class","to","train","the","reward","model","Args",":","config","(","ConfigModel",")",":","Config","parameters","for","the","model","Attributes",":","model","(","RewardModel",")",":","Reward","model","config","(","ConfigModel",")",":","Config","parameters","for","the","model","optimizer","(","torch.optim",")",":","Optimizer","for","the","model","loss_function","(","torch.nn",")",":","Loss","function","for","the","model","validation_flag","(","bool",")",":","Flag","to","indicate","if","the","validation","dataset","is","available","train_dataset","(","RewardDataset",")",":","Dataset","for","training","validation_dataset","(","RewardDataset",")",":","Dataset","for","validation","train_dataloader","(","DataLoader",")",":","Dataloader","for","training","validation_dataloader","(","DataLoader",")",":","Dataloader","for","validation","scheduler","(","torch.optim.lr_scheduler",")",":","Scheduler","for","the","optimizer","training_stats","(","List","[","Dict","]",")",":","List","of","dictionaries","with","the","training","statistics","model_engine","(","ModelEngine",")",":","Model","engine","to","train","the","model","using","deepspeed","accelerator","(","Accelerator",")",":","Accelerator","to","train","the","model","using","accelerate","by","HF",".","Methods",":","train",":","Train","the","reward","model","save_checkpoints",":","Save","the","checkpoints","of","the","model","load_checkpoints",":","Load","the","checkpoints","of","the","model"]
docstr: Save the checkpoints of the model
  
  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","checkpoints","of","the","model","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 checkpoints of the model
  
      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} ..."
      )
  
      # get the 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)
  
      # save the checkpoint
      if self.config.deepspeed_enable:
          client_state = {
              "epoch": current_epoch,
              "step": current_step,
          }
          self.model_engine.save_checkpoint(
              path, client_state=client_state
          )
      else:
          torch.save(
              {
                  "state_dict": self.reward.model.state_dict(),
                  "optim_state_dict": self.optimizer.state_dict(),
                  "scheduler_state_dict": self.scheduler.state_dict(),
                  "training_stats": self.training_stats,
                  "epoch": current_epoch,
                  "step": current_step,
              },
              path,
          )
  
imports: import json
  import shutil
  import os
  import deepspeed
  import torch
  from accelerate import Accelerator
  from beartype import beartype
  from beartype.typing import Iterable, Tuple
  from einops.layers.torch import Rearrange
  from torch.utils.data import Dataset, DataLoader
  from transformers import AutoModel, AutoTokenizer
  from chatllama.rlhf.config import ConfigReward
  from chatllama.rlhf.model_list import 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 RewardTrainer obj's underlying member method code with example usage: obj.save_checkpoint(current_epoch, current_step, max_epochs, max_steps) without return types