def METHOD_NAME(self, stage, stage_loss, epoch=None):
    """Gets called at the end of a stage."""
    if stage == sb.Stage.TRAIN:
        self.train_loss = stage_loss
    elif stage == sb.Stage.VALID:
        print("Completed epoch %d" % epoch)
        print("Train SI-SNR: %.3f" % -self.train_loss)
        print("Valid SI-SNR: %.3f" % -stage_loss)
    elif stage == sb.Stage.TEST:
        print("Test SI-SNR: %.3f" % -stage_loss)