Correctly restore TFA option when the function raises an Exception
@@ -334,7 +334,9 @@ def disable_tfa_custom_ops(func):
def _wrapper(*args, **kwargs):
previous_value = tfa.options.TF_ADDONS_PY_OPS
tfa.options.TF_ADDONS_PY_OPS = True
+ try:
outputs = func(*args, **kwargs)
+ finally:
tfa.options.TF_ADDONS_PY_OPS = previous_value
return outputs
