diff options
Diffstat (limited to 'tensorflow/python/ops/while_v2.py')
| -rw-r--r-- | tensorflow/python/ops/while_v2.py | 4 |
1 files changed, 4 insertions, 0 deletions
diff --git a/tensorflow/python/ops/while_v2.py b/tensorflow/python/ops/while_v2.py index 875be31602..6791e1cd61 100644 --- a/tensorflow/python/ops/while_v2.py +++ b/tensorflow/python/ops/while_v2.py @@ -24,6 +24,7 @@ from __future__ import division from __future__ import print_function import collections +import sys from tensorflow.core.framework import attr_value_pb2 from tensorflow.python.eager import function @@ -33,6 +34,7 @@ from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape from tensorflow.python.ops import array_ops from tensorflow.python.ops import cond_v2_impl as cond_v2 +from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import control_flow_util from tensorflow.python.ops import gen_functional_ops from tensorflow.python.ops import gradients_impl @@ -41,6 +43,8 @@ from tensorflow.python.util import nest # pylint: disable=protected-access +control_flow_ops._while_v2 = sys.modules[__name__] + # TODO(b/79881896): Handle external control dependencies. tf.while_loop allows # control dependencies on external nodes with at least 1 output. # Another idea is to create const nodes outside the loop and add control edges |