fix(transforms): restore tracing state after exceptions

Signed-off-by: kyinhub <kevinpyin@gmail.com>

Author: yinkev

monai/transforms/inverse.py

@@ -405,8 +405,10 @@ def trace_transform(self, to_trace: bool):
         """Temporarily set the tracing status of a transform with a context manager."""
         prev = self.tracing
         self.tracing = to_trace
-        yield
-        self.tracing = prev
+        try:
+            yield
+        finally:
+            self.tracing = prev
 
 
 class InvertibleTransform(TraceableTransform, InvertibleTrait):

tests/transforms/inverse/test_traceable_transform.py

@@ -29,6 +29,17 @@ def pop(self, data):
 
 class TestTraceable(unittest.TestCase):
 
+    def test_trace_transform_restores_state_after_exception(self):
+        transform = _TraceTest()
+        transform.tracing = True
+
+        with self.assertRaisesRegex(RuntimeError, "expected failure"):
+            with transform.trace_transform(False):
+                self.assertFalse(transform.tracing)
+                raise RuntimeError("expected failure")
+
+        self.assertTrue(transform.tracing)
+
     def test_default(self):
         expected_key = "_transforms"
         a = _TraceTest()