1 files changed, 7 insertions, 0 deletions

diff --git a/training/functional.py b/training/functional.py
index 83e70e2..62b8260 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -157,6 +157,7 @@ def save_samples(
                 width=image_size,
                 generator=gen,
                 guidance_scale=guidance_scale,
+                sag_scale=0,
                 num_inference_steps=num_steps,
                 output_type='pil'
             ).images
@@ -273,6 +274,12 @@ def loss_step(
         layout=latents.layout,
         device=latents.device,
         generator=generator
+    ) + 0.1 * torch.randn(
+        latents.shape[0], latents.shape[1], 1, 1,
+        dtype=latents.dtype,
+        layout=latents.layout,
+        device=latents.device,
+        generator=generator
     )
     bsz = latents.shape[0]
     # Sample a random timestep for each image