1 files changed, 3 insertions, 3 deletions

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