Skip to content

Commit 14b9754

Browse files
authored
[train_unconditional] fix applying clip_grad_norm_ (#721)
fix clip_grad_norm_
1 parent 6b22192 commit 14b9754

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

examples/unconditional_image_generation/train_unconditional.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -143,7 +143,8 @@ def transforms(examples):
143143
loss = F.mse_loss(noise_pred, noise)
144144
accelerator.backward(loss)
145145

146-
accelerator.clip_grad_norm_(model.parameters(), 1.0)
146+
if accelerator.sync_gradients:
147+
accelerator.clip_grad_norm_(model.parameters(), 1.0)
147148
optimizer.step()
148149
lr_scheduler.step()
149150
if args.use_ema:

0 commit comments

Comments
 (0)