From ae0f791a5e24457b4b0acdbd796a7dcce8ebd395 Mon Sep 17 00:00:00 2001 From: Will Cromar Date: Thu, 16 May 2024 21:55:28 +0000 Subject: [PATCH] improve test condition --- test/test_devices.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/test/test_devices.py b/test/test_devices.py index e978db3f47e..4588940fb80 100644 --- 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)