diff options
Diffstat (limited to 'tensorflow/python/keras/metrics.py')
| -rw-r--r-- | tensorflow/python/keras/metrics.py | 4 |
1 files changed, 2 insertions, 2 deletions
diff --git a/tensorflow/python/keras/metrics.py b/tensorflow/python/keras/metrics.py index 2dde9ee41f..9b87170ebe 100644 --- a/tensorflow/python/keras/metrics.py +++ b/tensorflow/python/keras/metrics.py @@ -55,7 +55,7 @@ from tensorflow.python.ops import nn from tensorflow.python.ops import state_ops from tensorflow.python.ops import variable_scope as vs from tensorflow.python.ops import weights_broadcast_ops -from tensorflow.python.training import distribute as distribute_lib +from tensorflow.python.training import distribution_strategy_context from tensorflow.python.util import tf_decorator from tensorflow.python.util.tf_export import tf_export @@ -111,7 +111,7 @@ def result_wrapper(result_fn): def decorated(metric_obj, *args): """Decorated function with merge_call.""" - tower_context = distribute_lib.get_tower_context() + tower_context = distribution_strategy_context.get_tower_context() if tower_context is None: # if in cross tower context already result_t = result_fn(*args) else: |