Reenable AMP on XLA (#61861)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/61861
Fixes https://github.com/pytorch/pytorch/issues/61804
Signed-off-by: Edward Z. Yang <ezyang@fb.com>
Test Plan: Imported from OSS
Reviewed By: jbschlosser
Differential Revision: D29881903
Pulled By: ezyang
fbshipit-source-id: 91530c10fa37715bec33f477285da119415a9da9
