Ensure the computation stays in single precision

[certik]

On my computer this change makes picoGPT several times faster.

Before:

$ time python gpt2.py 'Alan Turing theorized that computers would one day become' --models_dir ../gpt2/models 
generating: 100%|███████████████████████████████| 40/40 [00:16<00:00,  2.40it/s]
 the most powerful machines on the planet.

The computer is a machine that can perform complex calculations, and it can perform these calculations in a way that is very similar to the human brain.

python gpt2.py 'Alan Turing theorized that computers would one day become'    95.80s user 1.85s system 516% cpu 18.913 total

After:

$ time python gpt2.py 'Alan Turing theorized that computers would one day become' --models_dir ../gpt2/models
generating: 100%|███████████████████████████████| 40/40 [00:03<00:00, 12.11it/s]
 the most powerful machines on the planet.

The computer is a machine that can perform complex calculations, and it can perform these calculations in a way that is very similar to the human brain.

python gpt2.py 'Alan Turing theorized that computers would one day become'    22.94s user 0.72s system 428% cpu 5.526 total

The inference went from 16s to 3s.