improve test condition

pytorch · May 16, 2024 · ae0f791 · ae0f791
1 parent 02d8f9f
commit ae0f791
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/test/test_devices.py b/test/test_devices.py
@@ -50,12 +50,10 @@ def test_step(self):
     self.assertEqual(met.counter_value('MarkStep'), 1)
 
   def test_step_exception(self):
-    try:
+    with self.assertRaisesRegex(RuntimeError, 'Expected error'):
       with xla.step():
         torch.ones((3, 3), device=xla.device())
-        raise RuntimeError("Expected error")
-    except RuntimeError:
-      pass
+        raise RuntimeError('Expected error')
 
     self.assertEqual(met.counter_value('MarkStep'), 1)