xla

05ca2241 - Revert "fix batch_norm amp autocast" (#8547)