From 14b97549230ce4a40c0775e075a11c3891dc0446 Mon Sep 17 00:00:00 2001 From: Suraj Patil Date: Tue, 4 Oct 2022 19:04:05 +0200 Subject: [PATCH] [train_unconditional] fix applying clip_grad_norm_ (#721) fix clip_grad_norm_ --- examples/unconditional_image_generation/train_unconditional.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/unconditional_image_generation/train_unconditional.py b/examples/unconditional_image_generation/train_unconditional.py index 243e433a5b..6955c9f60a 100644 --- 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: