Allow individual models to scale the loss after it is calculated.
This commit is contained in:
1 parent
f743ccf7ef
commit
d14f6e567a
3 files changed
+11
No files matched your search
@@ -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
|
||||
@@ -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
|
||||
Reference in new issue
Block a user