Internal change

PiperOrigin-RevId: 321057047

Internal change
PiperOrigin-RevId: 321057047
b0693846 · Ken Franko · A. Unique TensorFlower · be8c2556 · b0693846
Commit b0693846 authored Jul 13, 2020 by Ken Franko Committed by A. Unique TensorFlower Jul 13, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

official/nlp/tasks/tagging.py official/nlp/tasks/tagging.py +1 -1

No files found.
--- a/official/nlp/tasks/tagging.py
+++ b/official/nlp/tasks/tagging.py
@@ -262,7 +262,7 @@ def predict(task: TaggingTask, params: cfg.DataConfig,
          label_mask=label_mask,
          sentence_ids=sentence_ids)

-    outputs = tf.distribute.get_strategy().experimental_run_v2(
+    outputs = tf.distribute.get_strategy().run(
        _replicated_step, args=(next(iterator),))
    return tf.nest.map_structure(
        tf.distribute.get_strategy().experimental_local_results, outputs)