Skip to content

Commit

Permalink
Merge pull request #470 from yelboudouri/master
Browse files Browse the repository at this point in the history
Unscale gradients before clipping
  • Loading branch information
milesial authored Feb 11, 2024
2 parents 2f62e6b + 52b4f14 commit de41eaa
Showing 1 changed file with 1 addition and 0 deletions.
1 change: 1 addition & 0 deletions train.py
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,7 @@ def train_model(

optimizer.zero_grad(set_to_none=True)
grad_scaler.scale(loss).backward()
grad_scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), gradient_clipping)
grad_scaler.step(optimizer)
grad_scaler.update()
Expand Down

0 comments on commit de41eaa

Please sign in to comment.