Fix loss=nan

1 files changed, 2 insertions, 2 deletions

diff --git a/training/functional.py b/training/functional.py
index 2d6553a..1baf9c6 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -314,7 +314,7 @@ def loss_step(
         target = noise_scheduler.get_velocity(latents, noise, timesteps)
         snr = target / (1 - target)
         snr /= snr + 1
-        snr_weights = torch.minimum(snr, torch.tensor([5], device=latents.device))
+        snr_weights = snr.fmin(torch.tensor([5], device=latents.device))
     else:
         raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")
 
@@ -334,7 +334,7 @@ def loss_step(
     else:
         loss = F.mse_loss(model_pred.float(), target.float(), reduction="none")
 
-    loss = (snr_weights * loss).mean()
+    loss = (snr_weights * loss).mean([1, 2, 3]).mean()
     acc = (model_pred == target).float().mean()
 
     return loss, acc, bsz