Merge pull request #32 from lopuhin/fix-tpu-training

Fix training on the TPU in Colab

tf/tpu_estimator.py

@@ -3083,6 +3083,12 @@ class _CapturingContext(control_flow_ops.ControlFlowContext):
         raise ValueError('{}: Op {} depends on TPU computation {}, '
                          'which is not allowed.'.format(self._message, op, c))
 
+  def to_control_flow_context_def(self, context_def, export_scope=None):
+    # pylint: disable=useless-super-delegation
+    # NOTE(slebedev): the method is required by `ControlFlowContext`.
+    super(_CapturingContext, self).to_control_flow_context_def(
+      context_def, export_scope)
+
   def __enter__(self):
     # pylint: disable=protected-access
     self._g = ops.get_default_graph()

tf/train.py

@@ -90,10 +90,6 @@ flags.DEFINE_integer("max_eval_batch", default=-1,
       help="Set -1 to turn off. Only used in test mode.")
 flags.DEFINE_bool("do_eval_only", default=False,
       help="Run evaluation only.")
-flags.DEFINE_string("eval_ckpt_path", None,
-      help="Checkpoint path for do_test evaluation."
-           "If set, model_dir will be ignored."
-           "If unset, will use the latest ckpt in model_dir.")
 flags.DEFINE_integer("start_eval_steps", default=10000,
       help="Which checkpoint to start with in `do_eval_only` mode.")
 flags.DEFINE_string("eval_split", "valid",