1 files changed, 5 insertions, 2 deletions

diff --git a/training/functional.py b/training/functional.py
index b9574ec..15b95ba 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -314,9 +314,12 @@ def loss_step(
     elif noise_scheduler.config.prediction_type == "v_prediction":
         target = noise_scheduler.get_velocity(latents, noise, timesteps)
 
+        p2_gamma = 1
+        p2_k = 1
+
         alpha_t = noise_scheduler.alphas_cumprod[timesteps].float()
-        snr = alpha_t / (1 - alpha_t)
-        snr_weights = snr.fmin(torch.tensor([5], device=latents.device)) / (snr + 1)
+        snr = 1.0 / (1 - alpha_t) - 1
+        snr_weights = 1 / (p2_k + snr) ** p2_gamma
         snr_weights = snr_weights[..., None, None, None]
     else:
         raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")