class_docstr: The Brain implementation for Tacotron2
docstr: Gets called at the end of an epoch
  
  Arguments
  ---------
  stage : sb.Stage
      One of sb.Stage.TRAIN, sb.Stage.VALID, sb.Stage.TEST
  stage_loss : float
      The average loss for all of the data processed in this stage.
  epoch : int
      The currently-starting epoch. This is passed
      `None` during the test stage.
docstr_tok: ["Gets","called","at","the","end","of","an","epoch","Arguments","--","--","--","--","-","stage",":","sb.Stage","One","of","sb.Stage.TRAIN",",","sb.Stage.VALID",",","sb.Stage.TEST","stage_loss",":","float","The","average","loss","for","all","of","the","data","processed","in","this","stage",".","epoch",":","int","The","currently-starting","epoch",".","This","is","passed","`","None","`","during","the","test","stage","."]
code: def on_stage_end(self, stage, stage_loss, epoch):
      """Gets called at the end of an epoch
  
      Arguments
      ---------
      stage : sb.Stage
          One of sb.Stage.TRAIN, sb.Stage.VALID, sb.Stage.TEST
      stage_loss : float
          The average loss for all of the data processed in this stage.
      epoch : int
          The currently-starting epoch. This is passed
          `None` during the test stage.
      """
  
      # Logs training samples every 10 epochs
      if stage == sb.Stage.TRAIN and (
          self.hparams.epoch_counter.current % 10 == 0
      ):
          if self.last_batch is None:
              return
  
          train_sample_path = os.path.join(
              self.hparams.progress_sample_path,
              str(self.hparams.epoch_counter.current),
          )
          if not os.path.exists(train_sample_path):
              os.makedirs(train_sample_path)
  
          (
              _,
              targets,
              _,
              labels,
              wavs,
              spk_embs,
              spk_ids,
          ) = self.last_batch
  
          train_sample_text = os.path.join(
              self.hparams.progress_sample_path,
              str(self.hparams.epoch_counter.current),
              "train_input_text.txt",
          )
          with open(train_sample_text, "w") as f:
              f.write(labels[0])
  
          train_input_audio = os.path.join(
              self.hparams.progress_sample_path,
              str(self.hparams.epoch_counter.current),
              "train_input_audio.wav",
          )
          torchaudio.save(
              train_input_audio,
              sb.dataio.dataio.read_audio(wavs[0]).unsqueeze(0),
              self.hparams.sample_rate,
          )
  
          _, mel_out_postnet, _, _, pred_mel_lengths = self.last_preds
  
          if self.hparams.log_audio_samples:
              waveform_ss = self.vocoder.decode_batch(
                  mel_out_postnet[0]
              )
              train_sample_audio = os.path.join(
                  self.hparams.progress_sample_path,
                  str(self.hparams.epoch_counter.current),
                  "train_output_audio.wav",
              )
              torchaudio.save(
                  train_sample_audio,
                  waveform_ss.squeeze(1).cpu(),
                  self.hparams.sample_rate,
              )
  
          if self.hparams.use_tensorboard:
              self.tensorboard_logger.log_audio(
                  f"{stage}/train_audio_target",
                  sb.dataio.dataio.read_audio(wavs[0]).unsqueeze(0),
                  self.hparams.sample_rate,
              )
              if self.hparams.log_audio_samples:
                  self.tensorboard_logger.log_audio(
                      f"{stage}/train_audio_pred",
                      waveform_ss.squeeze(1),
                      self.hparams.sample_rate,
                  )
              try:
                  self.tensorboard_logger.log_figure(
                      f"{stage}/train_mel_target", targets[0][0]
                  )
                  self.tensorboard_logger.log_figure(
                      f"{stage}/train_mel_pred", mel_out_postnet[0]
                  )
              except Exception:
                  # This is to avoid the code from crashing in case of a mel-spectrogram with one frame
                  pass
  
      # At the end of validation, we can write
      if stage == sb.Stage.VALID:
          # Update learning rate
          lr = self.optimizer.param_groups[-1]["lr"]
          self.last_epoch = epoch
  
          # The train_logger writes a summary to stdout and to the logfile.
          self.hparams.train_logger.log_stats(  # 1#2#
              stats_meta={"Epoch": epoch, "lr": lr},
              train_stats=self.last_loss_stats[sb.Stage.TRAIN],
              valid_stats=self.last_loss_stats[sb.Stage.VALID],
          )
  
          # The tensorboard_logger writes a summary to stdout and to the logfile.
          if self.hparams.use_tensorboard:
              self.tensorboard_logger.log_stats(
                  stats_meta={"Epoch": epoch, "lr": lr},
                  train_stats=self.last_loss_stats[sb.Stage.TRAIN],
                  valid_stats=self.last_loss_stats[sb.Stage.VALID],
              )
  
          # Save the current checkpoint and delete previous checkpoints.
          epoch_metadata = {
              **{"epoch": epoch},
              **self.last_loss_stats[sb.Stage.VALID],
          }
          self.checkpointer.save_and_keep_only(
              meta=epoch_metadata,
              min_keys=["loss"],
              ckpt_predicate=(
                  lambda ckpt: (
                      ckpt.meta["epoch"]
                      % self.hparams.keep_checkpoint_interval
                      != 0
                  )
              )
              if self.hparams.keep_checkpoint_interval is not None
              else None,
          )
          output_progress_sample = (
              self.hparams.progress_samples
              and epoch % self.hparams.progress_samples_interval == 0
          )
          if output_progress_sample:
              self.run_inference_sample(sb.Stage.VALID)
              self.hparams.progress_sample_logger.save(epoch)
  
      # We also write statistics about test data to stdout and to the logfile.
      if stage == sb.Stage.TEST:
          self.hparams.train_logger.log_stats(
              {"Epoch loaded": self.hparams.epoch_counter.current},
              test_stats=self.last_loss_stats[sb.Stage.TEST],
          )
          if self.hparams.use_tensorboard:
              self.tensorboard_logger.log_stats(
                  {"Epoch loaded": self.hparams.epoch_counter.current},
                  test_stats=self.last_loss_stats[sb.Stage.TEST],
              )
          if self.hparams.progress_samples:
              self.run_inference_sample(sb.Stage.TEST)
              self.hparams.progress_sample_logger.save("test")
  
imports: import torch
  import speechbrain
  import sys
  import logging
  from hyperpyyaml import load_hyperpyyaml
  from speechbrain.utils.text_to_sequence import text_to_sequence
  from speechbrain.utils.data_utils import scalarize
  import os
  from speechbrain.inference.vocoders import HIFIGAN
  import torchaudio
label_desc: Use image node_id 9 for calling the Tacotron2Brain obj's underlying member method code with example usage: obj.on_stage_end(stage, stage_loss, epoch) without return types