Gradient clipping: norm is always 2, device might be undefined.
🐛 Bug
I just re-skimmed the gradient clipping operation and found some inconsistency:
(1) device is only defined for non-'inf' norms, see here. While fixing this, I realized that
(2) norm_type is always 2.
So the two options I see are:
(1) Keep the existing behavior and remove the inf branch --> always apply 2-norm.
(2) Expose norm_type as a trainer argument.
Which option do you prefer?
I've also assumed that test_grad_norm.py tests gradient_clipping but it actually tests gradient norm tracking, so
I'd add a corresponding test case to the PR to fix the issue above ;)
(Skipping reproducibility and version info as this is easy to observe from looking at the code :))