More tolerant control flow grad code. It currently fails sometimes inside tf.function

PiperOrigin-RevId: 314213428 Change-Id: I0b0d1ae88401016d3b617cf843192708bda3e695
2020-06-01 15:17:02 -07:00 · 2020-06-01 15:17:02 -07:00 · 926090c4d4
commit 926090c4d4
parent c98e236060
1 changed files with 2 additions and 0 deletions
--- a/tensorflow/python/ops/control_flow_grad.py
+++ b/tensorflow/python/ops/control_flow_grad.py
@ -210,6 +210,8 @@ def _EnterGrad(op, grad):
  # pylint: disable=protected-access
  grad_ctxt = graph._get_control_flow_context()
  # pylint: enable=protected-access
+  if grad_ctxt is None:
+    return grad
  if not grad_ctxt.back_prop:
    # Skip gradient computation, if the attribute `back_prop` is false.
    return grad