Allow individual models to scale the loss after it is calculated.

2026-06-25 10:13:58 -06:00
parent f743ccf7ef
commit d14f6e567a
3 changed files with 11 additions and 0 deletions
--- a/extensions_built_in/sd_trainer/SDTrainer.py
+++ b/extensions_built_in/sd_trainer/SDTrainer.py
@@ -827,6 +827,9 @@ class SDTrainer(BaseSDTrainProcess):
                loss = torch.nn.functional.mse_loss(pred.float(), target.float(), reduction="none")
            
            loss = loss * local_loss_scale
+            
+            # apply model specific loss scaling
+            loss = self.sd.scale_loss(loss)
                
            do_weighted_timesteps = False
            if self.sd.is_flow_matching: