TPUEstimator.export_savedmodel() saves a SavedModel with both TPU and CPU graphs.

PiperOrigin-RevId: 198229550

diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py
index 4465833..c8c08a5 100644 (file)
--- a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py
+++ b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py
@@ -1807,7 +1807,7 @@ class TPUEstimator(estimator_lib.Estimator):
       export_outputs['classes'] =
         export_output_lib.ClassificationOutput(classes=classes)
 
-    tpu.outside_compilation(host_call, logits)
+    tpu.outside_compilation(host_call, [logits])
 
     ...
   ```
@@ -1969,7 +1969,7 @@ class TPUEstimator(estimator_lib.Estimator):
                              input_receiver_fn_map[mode]}
     export_tags = [tag_constants.SERVING, tag_constants.TPU]
     mode = _REWRITE_FOR_INFERENCE_MODE
-    try:
+    if self._export_to_tpu:
       (super(TPUEstimator, self).
        _add_meta_graph_for_mode(builder,
                                 input_receiver_fn_map,
@@ -1978,9 +1978,6 @@ class TPUEstimator(estimator_lib.Estimator):
                                 save_variables=False,
                                 mode=mode,
                                 export_tags=export_tags))
-    except Exception as error:  # pylint: disable=broad-except
-      logging.warning('Saving meta graph for TPU failed: {}.'
-                      .format(str(error)))
 
   def _call_model_fn(self, features, labels, mode, config):
     if mode == _REWRITE_FOR_INFERENCE_MODE: