could not run example with disable_eager_execution()

hi ALL,
I'm new to tf2.x, and following the tutorials https://github.com/tensorflow/text/blob/a2feab1cb3f30dedaf57797faf74eaebed30f405/docs/tutorials/classify_text_with_bert.ipynb
to run bert in graph mode, but got errors after I add tf.compat.v1.disable_eager_execution() and remove code relevant to eager mode.

my tensorflow version is 2.5.1
the errors are

Traceback (most recent call last):
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/client/session.py", line 1375, in _do_call
    return fn(*args)
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/client/session.py", line 1358, in _run_fn
    self._extend_graph()
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/client/session.py", line 1398, in _extend_graph
    tf_session.ExtendSession(self._session)
tensorflow.python.framework.errors_impl.InvalidArgumentError: Node 'training/AdamWeightDecay/gradients/gradients/BERT_encoder/cond_grad/If': Connecting to invalid output 7 of source node BERT_encoder/cond which has 7 outputs. Try using tf.compat.v1.experimental.output_all_intermediates(True).

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "classify_text_with_bert.py", line 458, in <module>
    history = classifier_model.fit(x=train_ds,
  File "/usr/local/lib/python3.8/dist-packages/keras/engine/training_v1.py", line 777, in fit
    return func.fit(
  File "/usr/local/lib/python3.8/dist-packages/keras/engine/training_arrays_v1.py", line 640, in fit
    return fit_loop(
  File "/usr/local/lib/python3.8/dist-packages/keras/engine/training_arrays_v1.py", line 188, in model_iteration
    validation_steps = training_utils_v1.infer_steps_for_dataset(
  File "/usr/local/lib/python3.8/dist-packages/keras/engine/training_utils_v1.py", line 1672, in infer_steps_for_dataset
    size = backend.get_value(tf.data.experimental.cardinality(dataset))
  File "/usr/local/lib/python3.8/dist-packages/keras/backend.py", line 3696, in get_value
    return x.eval(session=get_session((x,)))
  File "/usr/local/lib/python3.8/dist-packages/keras/backend.py", line 711, in get_session
    _initialize_variables(session)
  File "/usr/local/lib/python3.8/dist-packages/keras/backend.py", line 1161, in _initialize_variables
    is_initialized = session.run(
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/client/session.py", line 967, in run
    result = self._run(None, fetches, feed_dict, options_ptr,
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/client/session.py", line 1190, in _run
    results = self._do_run(handle, final_targets, final_fetches,
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/client/session.py", line 1368, in _do_run
    return self._do_call(_run_fn, feeds, fetches, targets, options,
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/client/session.py", line 1394, in _do_call
    raise type(e)(node_def, op, message)  # pylint: disable=no-value-for-parameter
tensorflow.python.framework.errors_impl.InvalidArgumentError: Node 'training/AdamWeightDecay/gradients/gradients/BERT_encoder/cond_grad/If': Connecting to invalid output 7 of source node BERT_encoder/cond which has 7 outputs. Try using tf.compat.v1.experimental.output_all_intermediates(True).

anyone could help?

Reproduce the failure using the linked classify_text_with_bert.ipynb tutorial, TensorFlow 2.5.1, and tf.compat.v1.disable_eager_execution(). Start at classify_text_with_bert.py around line 458 and trace the graph-mode model.fit call and BERT_encoder/cond error. Done means the example has a documented or verified graph-mode path without the reported invalid-output failure.