enable autocast for xla (#48570)
Summary:
For enabling amp in torch/xla, see [this](https://github.com/pytorch/xla/pull/2654).
Pull Request resolved: https://github.com/pytorch/pytorch/pull/48570
Reviewed By: ezyang
Differential Revision: D26120627
Pulled By: ailzhang
fbshipit-source-id: 32627b17c04bfdad128624676ea9bf6f117bc97d
