[JIT][Autocast] Don't cast softmax on CPU
In eager autocasting, softmax is only up-casted on GPU (and not on CPU).
This fixes the JIT implementation to do the same.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/76661
Approved by: https://github.com/jjsjann123, https://github.com/eellison
