added cpu 16 bit

2026-09-13 12:50:26 +08:00 · 2019-07-24 19:05:20 -04:00
parent 65ce10c255
commit ed9d977c4a
2 changed files with 21 additions and 1 deletions
@@ -111,6 +111,26 @@ def test_cpu_model():
    run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)


+def test_cpu_model_with_amp():
+    """
+    Make sure model trains on CPU
+    :return:
+    """
+
+    trainer_options = dict(
+        progress_bar=False,
+        experiment=get_exp(),
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.4,
+        use_amp=True
+    )
+
+    model, hparams = get_model()
+
+    run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
+
+
 def test_all_features_cpu_model():
    """
    Test each of the trainer options