Low freq noise with randomized strength

1 files changed, 8 insertions, 1 deletions

diff --git a/training/functional.py b/training/functional.py
index ac89048..2d582bf 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -277,7 +277,14 @@ def loss_step(
     )
 
     if low_freq_noise != 0:
-        noise = noise * (1 - low_freq_noise) + low_freq_noise * torch.randn(
+        low_freq_factor = low_freq_noise * torch.randn(
+            latents.shape[0], 1, 1, 1,
+            dtype=latents.dtype,
+            layout=latents.layout,
+            device=latents.device,
+            generator=generator
+        )
+        noise = noise * (1 - low_freq_factor) + low_freq_factor * torch.randn(
             latents.shape[0], latents.shape[1], 1, 1,
             dtype=latents.dtype,
             layout=latents.layout,