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: Load the checkpoints of the model
  
  Returns:
      Tuple[int, int]: The current epoch and step
          from which you should resume the training
docstr_tok: ["Load","the","checkpoints","of","the","model","Returns",":","Tuple","[","int",",","int","]",":","The","current","epoch","and","step","from","which","you","should","resume","the","training"]
code: def load_checkpoint(
      self,
  ) -> Tuple[int, int]:
      """Load the checkpoints of the model
  
      Returns:
          Tuple[int, int]: The current epoch and step
              from which you should resume the training
      """
  
      print("Looking for checkpoints...")
      # look for the checkpoints
      path = ModelLoader.check_model_path(
          config=self.config,
          is_checkpoint=True,
          current_epoch=None,
      )
  
      # check if a checkpoint exists
      if path is not None:
          print("Loading ...")
  
          if self.config.deepspeed_enable:
              # try to load the checkpoint
              try:
                  _, client_state = self.model_engine.load_checkpoint(
                      path
                  )
              except Exception:
                  print(
                      "Checkpoint corrupted!"
                      "Try to remove the last checkpoint."
                      "Now Starting from epoch 0, step 0"
                  )
                  return 0, 0
              # load epoch and step to resume loops
              epoch = client_state["epoch"]
              step = client_state["step"]
          else:
              # try to load the checkpoint
              try:
                  checkpoint = torch.load(path)
              except Exception:
                  print(
                      "Checkpoint corrupted!"
                      "Try to remove the last checkpoint."
                      "Now Starting from epoch 0, step 0"
                  )
                  return 0, 0
  
              # load the model parameters and optimizer parameters
              # from the checkpoint
              epoch = checkpoint["epoch"]
              self.reward.model.load_state_dict(
                  checkpoint["state_dict"]
              )
              self.optimizer.load_state_dict(
                  checkpoint["optim_state_dict"]
              )
              self.scheduler.load_state_dict(
                  checkpoint["scheduler_state_dict"]
              )
              self.training_stats = checkpoint["training_stats"]
              step = checkpoint["step"]
          return epoch, step + 1  # return the next episode to train
      return 0, 0
  
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 3 for calling the RewardTrainer obj's underlying member method code with example usage: obj.load_checkpoint() and returns: int,  int, epoch,  unknown, int,  int, int,  int