When doing an audio loss, show the img and audio loss in the loss log

extensions_built_in/sd_trainer/SDTrainer.py

@@ -994,6 +994,8 @@ class SDTrainer(BaseSDTrainProcess):
         if batch.audio_pred is not None and batch.audio_target is not None:
             audio_loss = torch.nn.functional.mse_loss(batch.audio_pred.float(), batch.audio_target.float(), reduction="mean")
             audio_loss = audio_loss * self.train_config.audio_loss_multiplier
+            self.additional_logs['loss/img'] = loss.item()
+            self.additional_logs['loss/audio'] = audio_loss.item()
             loss = loss + audio_loss
 
         # check for additional losses