diff --git a/tensorflow/core/tpu/tpu_library_init_fns.inc b/tensorflow/core/tpu/tpu_library_init_fns.inc index f7f247596fe..b88705e4100 100644 --- a/tensorflow/core/tpu/tpu_library_init_fns.inc +++ b/tensorflow/core/tpu/tpu_library_init_fns.inc @@ -1,4 +1,5 @@ #include "third_party/tensorflow/stream_executor/tpu/tpu_executor_c_api.h" + namespace { tensorflow::Status SetTpuConfigStructFns(void* library_handle) { diff --git a/tensorflow/core/tpu/tpu_node_device.cc b/tensorflow/core/tpu/tpu_node_device.cc index 979bfab47db..42a1533a97c 100644 --- a/tensorflow/core/tpu/tpu_node_device.cc +++ b/tensorflow/core/tpu/tpu_node_device.cc @@ -90,6 +90,9 @@ Status TpuPaddedShapeFn(const Tensor& tensor, xla::Shape* shape) { ApiConverter::StackHelper<XLA_Shape> tpu_shape; tpu::ExecutorApiFn()->XlaShapeToTpuPaddedShapeFn( &se_shape.value, &tpu_shape.value, status.c_status); + if (!status.ok()) { + return status.status(); + } *shape = tpu_shape.AsCpp<xla::Shape>(); return Status::OK(); }