def learn(self):
        """
        Samples batches from `self.store`, updates model and periodically evaluates it on `self.eval_dataloader`
        """

        self.prepare_learning()

        tbar = tqdm(
            total=self.total_steps, disable=not self.accelerator.is_local_main_process
        )

        for _ in range(self.config.train.epochs):
            for _ in range(self.n_updates_per_batch):
                for batch in self.train_dataloader:
                    forward_time = time()
                    loss, stats = self.loss(batch)
                    forward_time = time() - forward_time

                    backward_time = time()
                    self.accelerator.backward(loss)
                    if self.config.train.max_grad_norm > 0 and self.accelerator.sync_gradients:
                        self.accelerator.clip_grad_norm_(self.model.parameters(), self.config.train.max_grad_norm)
                    backward_time = time() - backward_time

                    self.opt.step()
                    self.opt.zero_grad()
                    self.scheduler.step()
                    self.iter_count += 1

                    if self.iter_count % self.config.train.checkpoint_interval == 0:
                        self.save()

                    if self.iter_count % self.config.train.eval_interval == 0:
                        results = self.evaluate()

                        results.update(stats)
                        results.update(
                            {
                                "forward_time": forward_time,
                                "backward_time": backward_time,
                            }
                        )
                        self.accelerator.log(results)

                    desc = ", ".join(f"{k}: {v:.2f}" for k, v in stats.items())
                    tbar.set_description(desc)
                    tbar.update()

                    if self.iter_count >= self.total_steps:
                        self.save()
                        return self.evaluate()

                self.post_backward_callback()

            self.post_epoch_callback()

        Samples batches from `self.store`, updates model and periodically evaluates it on `self.eval_dataloader`
        
promptist/trlx/trlx/model/accelerate_base_model.py