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.
      """
  
      # Store the train loss until the validation stage.
      stage_stats = {"loss": stage_loss}
      stage_stats["PPL"] = math.exp(stage_loss)
      if stage == sb.Stage.TRAIN:
          self.train_stats = stage_stats
      else:
          stage_stats["BLEU_4"] = self.bleu_4_metric.summarize("BLEU")
          stage_stats["BLEU_2"] = self.bleu_2_metric.summarize("BLEU")
      # Perform end-of-iteration things, like annealing, logging, etc.
      if stage == sb.Stage.VALID:
          # Update learning rate
          old_lr, new_lr = self.hparams.lr_annealing(epoch)
          sb.nnet.schedulers.update_learning_rate(
              self.optimizer, new_lr
          )
  
          # The train_logger writes a summary to stdout and to the logfile.
  
          self.hparams.train_logger.log_stats(
              stats_meta={"epoch": epoch, "lr": old_lr},
              train_stats=self.train_stats,
              valid_stats=stage_stats,
          )
          # Save the current checkpoint and delete previous checkpoints.
          self.checkpointer.save_and_keep_only(
              meta={"PPL": stage_stats["PPL"]},
              min_keys=["PPL"],
          )
          if epoch == hparams["number_of_epochs"] - 1:
              with open(self.hparams.bleu_4_valid_file, "w") as w:
                  self.bleu_4_metric.write_stats(w)
                  for i in range(len(self.hyps)):
                      w.write(
                          "target: " + str(self.references[i]) + "\n"
                      )
                      w.write("predicted:" + str(self.hyps[i]) + "\n")
                      w.write(
                          "++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++\n"
                      )
  
      # We also write statistics about test data to stdout and to the logfile.
      elif stage == sb.Stage.TEST:
          self.hparams.train_logger.log_stats(
              stats_meta={
                  "Epoch loaded": self.hparams.epoch_counter.current
              },
              test_stats=stage_stats,
          )
          with open(self.hparams.bleu_4_test_file, "w") as w:
              self.bleu_4_metric.write_stats(w)
              for i in range(len(self.hyps)):
                  w.write("target: " + str(self.references[i]) + "\n")
                  w.write("predicted:" + str(self.hyps[i]) + "\n")
                  w.write(
                      "++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++\n"
                  )
  
file_path: speechbrain/recipes/MultiWOZ/response_generation/llama2/train_with_llama2.py
imports: import sys
  import speechbrain
  import torch
  from itertools import chain
  from hyperpyyaml import load_hyperpyyaml
  from speechbrain.utils.distributed import run_on_main
  import math
  from speechbrain.dataio.batch import PaddedBatch
label_desc: Use image node_id 5 for calling the ResGenBrain obj's underlying member method code with example usage: obj.on_stage_end(stage, stage_loss, epoch) without return types