diff options
Diffstat (limited to 'tensorflow/python/framework/function.py')
| -rw-r--r-- | tensorflow/python/framework/function.py | 6 |
1 files changed, 3 insertions, 3 deletions
diff --git a/tensorflow/python/framework/function.py b/tensorflow/python/framework/function.py index caa604999c..14d72d8a3d 100644 --- a/tensorflow/python/framework/function.py +++ b/tensorflow/python/framework/function.py @@ -489,10 +489,10 @@ class _DefinedFunction(object): # Adds this function into 'g'. # pylint: disable=protected-access - if context.in_graph_mode(): - g._add_function(self) - else: + if context.executing_eagerly(): context.context().add_function_def(self.definition) + else: + g._add_function(self) # pylint: enable=protected-access # Ensures related sub-routines are defined in 'g', too. |