Fix context manager not resetting dc mode in case of exception

apache · Feb 20, 2020 · 7e23173 · 7e23173
1 parent f35c20e
commit 7e23173
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/python/mxnet/_deferred_compute.py b/python/mxnet/_deferred_compute.py
@@ -54,8 +54,10 @@ def context(state=True):
     # other code unexpectedly, when used in concurrent code."
     # https://github.com/apache/incubator-mxnet/issues/17495#issuecomment-585461965
     val = set_deferred_compute(state)
-    yield
-    set_deferred_compute(val)
+    try:
+        yield
+    finally:
+        set_deferred_compute(val)
 
 
 def get_symbol(input_arrays, output_arrays, input_names=None, *, sym_cls=Symbol):