RomanNum3ral
/
ai-toolkit
mirror of https://github.com/ostris/ai-toolkit.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Check for finite vs isnan on loss before backpropigating. to catch infinity overflows

This commit is contained in:
Jaret Burkett 2026-08-06 09:18:56 -06:00
parent edacd406b3
commit b811636ae4
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
extensions_built_in/sd_trainer/SDTrainer.py
View File

@ -939,7 +939,7 @@ class SDTrainer(BaseSDTrainProcess):
prior_loss = torch.nn.functional.mse_loss(pred.float(), prior_pred.float(), reduction="none")
prior_loss = prior_loss * prior_mask_multiplier * self.train_config.inverted_mask_prior_multiplier
if torch.isnan(prior_loss).any():
if torch.isnan(prior_loss).any() or not torch.isfinite(prior_loss):
print_acc("Prior loss is nan")
prior_loss = None
else:
@ -2166,7 +2166,7 @@ class SDTrainer(BaseSDTrainProcess):
loss = loss + preservation_loss
# check if nan
if torch.isnan(loss):
if not torch.isfinite(loss):
print_acc("loss is nan")
loss = torch.zeros_like(loss).requires_grad_(True)