Fix mps ci

pytorch · Aug 5, 2024 · 5413327 · 5413327
1 parent 8be85b4
commit 5413327
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/ignite/engine/test_create_supervised.py b/tests/ignite/engine/test_create_supervised.py
@@ -316,7 +316,8 @@ def _test_create_supervised_evaluator(
             # This is broken in 1.6.0 but will be probably fixed with 1.7.0
             err_msg_1 = "Expected all tensors to be on the same device"
             err_msg_2 = "Placeholder storage has not been allocated on MPS device"
-            with pytest.raises(RuntimeError, match=f"({err_msg_1}|{err_msg_2})"):
+            err_msg_3 = "Tensor for argument weight is on cpu but expected on mps"
+            with pytest.raises(RuntimeError, match=f"({err_msg_1}|{err_msg_2}|{err_msg_3})"):
                 evaluator.run(data)