    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.train()
        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 = self._train_epoch()
            valid_total_loss = self._valid_epoch()
            self.accelerator.log(
                {
                    "Epoch/Train Total Loss": train_total_loss,
                    "Epoch/Valid Total Loss": valid_total_loss,
                },
                step=self.epoch,
            )

            # Update scheduler
            self.accelerator.wait_for_everyone()
            self.scheduler.step()

            # 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:
                path = os.path.join(
                    self.checkpoint_dir,
                    "epoch-{:04d}_step-{:07d}_loss-{:.6f}".format(
                        self.epoch, self.step, valid_total_loss
                    ),
                )
                self.accelerator.save_state(path)
                json.dump(
                    self.checkpoints_path,
                    open(os.path.join(path, "ckpts.json"), "w"),
                    ensure_ascii=False,
                    indent=4,
                )

            # Save eval audios
            self.accelerator.wait_for_everyone()
            if self.accelerator.is_main_process and run_eval:
                for i in range(len(self.valid_dataloader.dataset.eval_audios)):
                    if self.cfg.preprocess.use_frame_pitch:
                        eval_audio = self._inference(
                            self.valid_dataloader.dataset.eval_mels[i],
                            eval_pitch=self.valid_dataloader.dataset.eval_pitchs[i],
                            use_pitch=True,
                        )
                    else:
                        eval_audio = self._inference(
                            self.valid_dataloader.dataset.eval_mels[i]
                        )
                    path = os.path.join(
                        self.checkpoint_dir,
                        "epoch-{:04d}_step-{:07d}_loss-{:.6f}_eval_audio_{}.wav".format(
                            self.epoch,
                            self.step,
                            valid_total_loss,
                            self.valid_dataloader.dataset.eval_dataset_names[i],
                        ),
                    )
                    path_gt = os.path.join(
                        self.checkpoint_dir,
                        "epoch-{:04d}_step-{:07d}_loss-{:.6f}_eval_audio_{}_gt.wav".format(
                            self.epoch,
                            self.step,
                            valid_total_loss,
                            self.valid_dataloader.dataset.eval_dataset_names[i],
                        ),
                    )
                    save_audio(path, eval_audio, self.cfg.preprocess.sample_rate)
                    save_audio(
                        path_gt,
                        self.valid_dataloader.dataset.eval_audios[i],
                        self.cfg.preprocess.sample_rate,
                    )

            self.accelerator.wait_for_everyone()

            self.epoch += 1

        # Finish training
        self.accelerator.wait_for_everyone()
        path = os.path.join(
            self.checkpoint_dir,
            "epoch-{:04d}_step-{:07d}_loss-{:.6f}".format(
                self.epoch, self.step, valid_total_loss
            ),
        )
        self.accelerator.save_state(path)
Training process
models/vocoders/diffusion/diffusion_vocoder_trainer.py