xla

7b538aac - prevent amp from running on TPU (#2921)

prevent amp from running on TPU (#2921) * prvent amp from running on TPU * Disable AMP for non-GPU device * allow cpu to run amp for testing purpose