fix precision problem?

anime-style · Feb 12, 2023 · afbcea1 · afbcea1
1 parent 9380060
commit afbcea1
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/lora_diffusion/cli_lora_pti.py b/lora_diffusion/cli_lora_pti.py
@@ -361,7 +361,9 @@ def loss_step(
 
         target = target * mask
 
-    loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
+    loss = F.mse_loss(model_pred.float(), target.float(), reduction="none")
+    loss = loss.sum(0).mean()
+
     return loss