Fix finite check on inverted masked prior

2026-08-07 09:49:12 -06:00
parent 685ce37a8d
commit 817f3dcbcb
1 changed files with 1 additions and 1 deletions
--- a/extensions_built_in/sd_trainer/SDTrainer.py
+++ b/extensions_built_in/sd_trainer/SDTrainer.py
@@ -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() or not torch.isfinite(prior_loss):
+            if not torch.isfinite(prior_loss).all():
                print_acc("Prior loss is nan")
                prior_loss = None
            else: