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/toolkit/models/base_model.py
+++ b/toolkit/models/base_model.py
@@ -1602,3 +1602,7 @@ class BaseModel:
    def get_model_to_train(self):
        # called to get model to attach LoRAs to. Can be overridden in child classes
        return self.unet
+    
+    def scale_loss(self, loss):
+        # called to get the loss scaler for the model. Can be overridden in child classes
+        return loss