🐛 Bug

The benchmarks in pytorch/xla does not work on CPU because of it is set to use AMP by default.

To Reproduce

Steps to reproduce the behavior:

1. Follow the instruction to run the benchmarks in the README.md.
2. Or run this command directly

python xla/benchmarks/experiment_runner.py --dynamo=openxla --xla=PJRT --test=eval --test=train --suite-name=torchbench --accelerator=cpu --output-dirname=/tmp/output --repeat=1 --print-subprocess --no-resume --dump-pytorch-xla-metrics

Expected behavior

Add a new arg to the CLI --amp. Let user configure it (disable AMP) when run on CPU.

Environment

• Reproducible on XLA backend [CPU/TPU/CUDA]: CPU
• torch_xla version: master or 2.7