From e64e03cb27b43f369c27c6629f1238103fa875e0 Mon Sep 17 00:00:00 2001 From: Leonard Lausen Date: Wed, 19 Feb 2020 21:11:51 +0000 Subject: [PATCH] Fix context manager not resetting dc mode in case of exception --- python/mxnet/_deferred_compute.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/python/mxnet/_deferred_compute.py b/python/mxnet/_deferred_compute.py index a280769060c6..1f8c6164c716 100644 --- a/python/mxnet/_deferred_compute.py +++ b/python/mxnet/_deferred_compute.py @@ -48,14 +48,17 @@ def set_deferred_compute(state): @contextlib.contextmanager def context(state=True): + """Set deferred compute state to `state` within context. Reset afterwards to previous value.""" # Like other MXNet context manager, this bleeds state across concurrent # code: "Context managers that have state should use Context Variables # instead of threading.local() to prevent their state from bleeding to # 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):