Merge pull request #470 from yelboudouri/master

Unscale gradients before clipping

train.py

@@ -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()