Fix #3954 - GPT2 is not traceable
This PR replaces the `math.sqrt(...)` computation with `(...)**0.5` which allows the model to be correctly traced by `torch.jit.trace`.
Need to remove the `math` import to pass the Flake test.<|||||>Just deleted the math import to make the code quality check happy - Thanks a lot for the PR @jazzcook15 and for linking the issue @minimaxir 