[train_unconditional] fix applying clip_grad_norm_ (#721)

fix clip_grad_norm_

[train_unconditional] fix applying clip_grad_norm_ (#721)
fix clip_grad_norm_
14b97549 · Suraj Patil · GitHub · 6b221920 · 14b97549
Unverified Commit 14b97549 authored Oct 04, 2022 by Suraj Patil Committed by GitHub Oct 04, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

examples/unconditional_image_generation/train_unconditional.py ...les/unconditional_image_generation/train_unconditional.py +2 -1

No files found.
--- a/examples/unconditional_image_generation/train_unconditional.py
+++ b/examples/unconditional_image_generation/train_unconditional.py
@@ -143,7 +143,8 @@ def main(args):
                loss = F.mse_loss(noise_pred, noise)
                accelerator.backward(loss)
-                accelerator.clip_grad_norm_(model.parameters(), 1.0)
+                if accelerator.sync_gradients:
+                    accelerator.clip_grad_norm_(model.parameters(), 1.0)
                optimizer.step()
                lr_scheduler.step()
                if args.use_ema: