grad_norm (line 32) should be before optimizer.step() (line 35)

[maximilianigl]

Or am I wrong?

[emited]

You are right again ! Don't hesitate if you want to send in a PR to correct this.

Cheers