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

This commit is contained in:
Jaret Burkett committed 2026-06-25 10:13:58 -06:00
1 parent f743ccf7ef
commit d14f6e567a
3 files changed
+11

No files matched your search

+4
View File
@@ -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
+4
View File
@@ -3161,3 +3161,7 @@ class StableDiffusion:
def get_model_to_train(self):
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