From 4de282cf8a80aaaf0879bfaffb259832fe41bcee Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Fri, 1 Mar 2019 18:56:04 -0800 Subject: [PATCH] Fix breakage related to HParam in tpu_estimator.py PiperOrigin-RevId: 236415795 --- tensorflow/python/tpu/tpu_estimator.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorflow/python/tpu/tpu_estimator.py b/tensorflow/python/tpu/tpu_estimator.py index 641ca281c15..38d85320472 100644 --- a/tensorflow/python/tpu/tpu_estimator.py +++ b/tensorflow/python/tpu/tpu_estimator.py @@ -2768,10 +2768,10 @@ class TPUEstimator(estimator_lib.Estimator): if is_export_mode: if mode == _REWRITE_FOR_INFERENCE_MODE: - params['use_tpu'] = True + _add_item_to_params(params, _USE_TPU_KEY, True) mode = model_fn_lib.ModeKeys.PREDICT else: - params['use_tpu'] = False + _add_item_to_params(params, _USE_TPU_KEY, False) with self._ctx.with_mode(mode) as ctx: model_fn_wrapper = _ModelFnWrapper(model_fn, config, params, ctx)