[train_unconditional] fix applying clip_grad_norm_ (huggingface#721)

fix clip_grad_norm_
Mistobaan · Oct 4, 2022 · 14b9754 · 14b9754
1 parent 6b22192
commit 14b9754
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/examples/unconditional_image_generation/train_unconditional.py b/examples/unconditional_image_generation/train_unconditional.py
@@ -143,7 +143,8 @@ def transforms(examples):
                 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: