From 4e50535478d59a6e418c4e153c6daa908ad240c5 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Sun, 28 Jun 2026 09:49:36 -0600 Subject: [PATCH] Allow models to provide an additional loss --- extensions_built_in/sd_trainer/SDTrainer.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/extensions_built_in/sd_trainer/SDTrainer.py b/extensions_built_in/sd_trainer/SDTrainer.py index b392246..6d19c64 100644 --- a/extensions_built_in/sd_trainer/SDTrainer.py +++ b/extensions_built_in/sd_trainer/SDTrainer.py @@ -944,6 +944,12 @@ class SDTrainer(BaseSDTrainProcess): loss = loss + additional_loss + if hasattr(self.sd, "get_additional_loss"): + additional_model_loss = self.sd.get_additional_loss(pred, target) + if additional_model_loss is not None: + loss = loss + additional_model_loss + self.additional_logs["additional_model_loss"] = additional_model_loss.item() + if self.train_config.max_loss_debug and self.train_config.max_loss is not None: if loss.item() > self.train_config.max_loss: print_acc(f"Loss {loss.item()} is greater than max loss {self.train_config.max_loss}. Clipping to max loss.")