message: Override params["use_tpu"] to be true when building TPU graph and to be false when building CPU/GPU graph in export_savedmodel().
diff: @@ -1681,6 +1681,9 @@ class _ModelFnWrapper(object):
  _add_item_to_params(params, _BATCH_SIZE_KEY, batch_size_for_model_fn)
  running_on_cpu = self._ctx.is_running_on_cpu(is_export_mode)
  + # In export mode, params['use_tpu'] has already been set based on mode
  + # (i.e. True for _REWRITE_FOR_INFERENCE_MODE, False otherwise).
  + if not is_export_mode:
  _add_item_to_params(params, _USE_TPU_KEY, not running_on_cpu)
  if not running_on_cpu:
  @@ -2524,8 +2527,8 @@ class TPUEstimator(estimator_lib.Estimator):
  """
  # We should only call model fn once and it should be inside `computation`
  # so that building the graph will happen under `rewrite_for_inference`.
  - mode = model_fn_lib.ModeKeys.PREDICT
  - estimator_spec = self._call_model_fn(features, labels, mode, config)
  + estimator_spec = super(TPUEstimator, self)._call_model_fn(
  + features, labels, mode, config)
  # We pick the TPU tensors out from `export_output` and later return them
  # from `computation` for rewriting.
  @@ -2752,8 +2755,6 @@ class TPUEstimator(estimator_lib.Estimator):
  def _model_fn(features, labels, mode, config, params):
  """A Estimator `model_fn` for TPUEstimator."""
  - with self._ctx.with_mode(mode) as ctx:
  - model_fn_wrapper = _ModelFnWrapper(model_fn, config, params, ctx)
  # `input_fn` is called in `train()`, `evaluate()`, and `predict()`,
  # but not in `export_savedmodel()`.
  @@ -2765,6 +2766,16 @@ class TPUEstimator(estimator_lib.Estimator):
  # Clear the bit.
  self._is_input_fn_invoked = None
  + if is_export_mode:
  + if mode == _REWRITE_FOR_INFERENCE_MODE:
  + params['use_tpu'] = True
  + mode = model_fn_lib.ModeKeys.PREDICT
  + else:
  + params['use_tpu'] = False
  +
  + with self._ctx.with_mode(mode) as ctx:
  + model_fn_wrapper = _ModelFnWrapper(model_fn, config, params, ctx)
  +
  # examples_hook is added to training_hooks for both CPU and TPU
  # execution.
  if self._log_every_n_steps is not None:
  