[dreambooth] fix applying clip_grad_norm_ (#686)

fix applying clip grad norm

[dreambooth] fix applying clip_grad_norm_ (#686)
fix applying clip grad norm
14f4af8f · Suraj Patil · GitHub · 2558977b · 14f4af8f
Unverified Commit 14f4af8f authored Oct 03, 2022 by Suraj Patil Committed by GitHub Oct 03, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

examples/dreambooth/train_dreambooth.py examples/dreambooth/train_dreambooth.py +2 -1

No files found.
--- a/examples/dreambooth/train_dreambooth.py
+++ b/examples/dreambooth/train_dreambooth.py
@@ -566,7 +566,8 @@ def main():
                    loss = F.mse_loss(noise_pred, noise, reduction="none").mean([1, 2, 3]).mean()

                accelerator.backward(loss)
-                accelerator.clip_grad_norm_(unet.parameters(), args.max_grad_norm)
+                if accelerator.sync_gradients:
+                    accelerator.clip_grad_norm_(unet.parameters(), args.max_grad_norm)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()