skip test_torch_adamw_equal for invalid parameter combinations

Change-Id: Ic3867ed78a636e88b884f304b68e897650a47ddc
microsoft · May 13, 2024 · 23625cf · 23625cf
1 parent cbef0bf
commit 23625cf
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/unit/ops/adam/test_cpu_adam.py b/tests/unit/ops/adam/test_cpu_adam.py
@@ -92,6 +92,8 @@ def test_fused_adam_equal(self, dtype, model_size):
 
     def test_torch_adamw_equal(self, dtype, model_size):
         if get_accelerator().is_available():
+            if dtype == torch.half:
+                pytest.skip("torch.optim.AdamW with half precision inf/nan output.")
             if ("amd" in pytest.cpu_vendor) and (dtype == torch.half):
                 pytest.skip("cpu-adam with half precision not supported on AMD CPUs")
             ref_param_device = get_accelerator().device_name()