Address bug in distributed strategies `Monitor` to allow running for >1 step.

PiperOrigin-RevId: 191667378

diff --git a/tensorflow/contrib/distribute/python/monitor.py b/tensorflow/contrib/distribute/python/monitor.py
index fe80bb4..7644ace 100644 (file)
--- a/tensorflow/contrib/distribute/python/monitor.py
+++ b/tensorflow/contrib/distribute/python/monitor.py
@@ -19,6 +19,7 @@ from __future__ import division
 from __future__ import print_function
 
 from tensorflow.python.eager import context
+from tensorflow.python.framework import errors
 from tensorflow.python.ops import variables
 
 
@@ -55,7 +56,9 @@ class Monitor(object):
 
   def run_steps(self, num_steps=None):
     step = 0
-    done = False
-    while done is not None and (num_steps is None or step < num_steps):
-      done = self._run_step()
-      step += 1
+    while num_steps is None or step < num_steps:
+      try:
+        self._run_step()
+        step += 1
+      except errors.OutOfRangeError:
+        break