Disable AMP by default on CPU

benchmarks/torchbench_model.py

@@ -373,6 +373,10 @@ def is_accelerator_tpu(self):
     return self.benchmark_experiment.accelerator == "tpu"
 
   def use_amp(self):
+    # AMP is only supported on cuda and tpu, not on cpu.
+    if self.benchmark_experiment.accelerator == "cpu":
+      logger.warning("AMP is not used due to running on CPU.")
+      return False
     return self.is_training() or self.model_name in config(
     ).dtype.force_amp_for_fp16_bf16_models
 

test/benchmarks/test_torchbench_model.py

@@ -0,0 +1,33 @@
+import unittest
+
+from benchmarks.torchbench_model import TorchBenchModel
+
+
+class MockExperiment:
+
+  def __init__(self, accelerator, test):
+    self.accelerator = accelerator
+    self.test = "train"
+
+
+class TorchBenchModelTest(unittest.TestCase):
+
+  def test_do_not_use_amp_on_cpu_and_warns(self):
+    experiment = MockExperiment("cpu", "train")
+    model = TorchBenchModel("torchbench or other", "super_deep_model",
+                            experiment)
+    with self.assertLogs('benchmarks.torchbench_model', level='WARNING') as cm:
+      use_amp = model.use_amp()
+      self.assertEqual(len(cm.output), 1)
+      self.assertIn("AMP is not used", cm.output[0])
+      self.assertFalse(use_amp)
+
+  def test_use_amp_on_cuda(self):
+    experiment = MockExperiment("cuda", "train")
+    model = TorchBenchModel("torchbench or other", "super_deep_model",
+                            experiment)
+    self.assertTrue(model.use_amp())
+
+
+if __name__ == '__main__':
+  unittest.main()