    def train_loop(self):
        """Training process"""
        self.accelerator.wait_for_everyone()

        # Dump config
        if self.accelerator.is_main_process:
            self._dump_cfg(self.config_save_path)
        _ = [self.model[key].train() for key in self.model]
        self.optimizer.zero_grad()

        # Sync and start training
        self.accelerator.wait_for_everyone()
        while self.epoch < self.max_epoch:
            self.logger.info("\n")
            self.logger.info("-" * 32)
            self.logger.info("Epoch {}: ".format(self.epoch))

            # Train and Validate
            train_total_loss, train_losses = self._train_epoch()
            for key, loss in train_losses.items():
                self.logger.info("  |- Train/{} Loss: {:.6f}".format(key, loss))
                self.accelerator.log(
                    {"Epoch/Train {} Loss".format(key): loss},
                    step=self.epoch,
                )
            self.accelerator.log(
                {
                    "Epoch/Train Total Loss": train_total_loss,
                },
                step=self.epoch,
            )

            # Update scheduler
            self.accelerator.wait_for_everyone()

            # Check save checkpoint interval
            run_eval = False
            if self.accelerator.is_main_process:
                save_checkpoint = False
                for i, num in enumerate(self.save_checkpoint_stride):
                    if self.epoch % num == 0:
                        save_checkpoint = True
                        run_eval |= self.run_eval[i]

            # Save checkpoints
            self.accelerator.wait_for_everyone()
            if self.accelerator.is_main_process and save_checkpoint:
                print("Saving..")
                state = {
                    "net": {key: self.model[key].state_dict() for key in self.model},
                    "optimizer": self.optimizer.state_dict(),
                    "scheduler": self.optimizer.scheduler_state_dict(),
                    "iters": self.step,
                    "epoch": self.epoch,
                }
                save_path = os.path.join(
                    self.checkpoint_dir,
                    "FAcodec_epoch_%05d_step_%05d.pth" % (self.epoch, self.iters),
                )
                torch.save(state, save_path)
                json.dump(
                    self.checkpoints_path,
                    open(os.path.join(self.checkpoint_dir, "ckpts.json"), "w"),
                    ensure_ascii=False,
                    indent=4,
                )

            self.accelerator.wait_for_everyone()

            self.epoch += 1

        # Finish training
        self.accelerator.wait_for_everyone()
        if self.accelerator.is_main_process:
            path = os.path.join(
                self.checkpoint_dir,
                "epoch-{:04d}_step-{:07d}".format(
                    self.epoch,
                    self.step,
                ),
            )
            print("Saving..")
            state = {
                "net": {key: self.model[key].state_dict() for key in self.model},
                "optimizer": self.optimizer.state_dict(),
                "scheduler": self.optimizer.scheduler_state_dict(),
                "iters": self.step,
                "epoch": self.epoch,
            }
            save_path = os.path.join(
                self.checkpoint_dir,
                "FAcodec_epoch_%05d_step_%05d.pth" % (self.epoch, self.iters),
            )
            torch.save(state, save_path)
Training process
models/codec/facodec/facodec_trainer.py