def save(self, directory=None):
        """Creates checkpoint of optimizer, scheduler and a model"""
        if self.accelerator.is_main_process:
            directory = directory or self.config.train.checkpoint_dir
            rank = self.accelerator.process_index
            directory = os.path.join(directory, "r%d-step%d" % (rank, self.iter_count))
            print("[I] Save states to %s" % directory)
            os.makedirs(directory, exist_ok=True)
            try:
                self.accelerator.save_state(directory)
            except OSError:
                print("[W] OSError while saving states...")
Creates checkpoint of optimizer, scheduler and a model
promptist/trlx/trlx/model/accelerate_base_model.py