From 2df63a4fa10b6da8c4f92bb55fe14bb30b0e4b73 Mon Sep 17 00:00:00 2001
From: Russell Power <power@google.com>
Date: Tue, 21 Jul 2020 17:48:06 -0700
Subject: [PATCH] Fix missing status check.
PiperOrigin-RevId: 322479633
Change-Id: I6452af8b6f106d6681307caab1c9b41e514d4694
---
tensorflow/core/tpu/tpu_library_init_fns.inc | 1 +
tensorflow/core/tpu/tpu_node_device.cc | 3 +++
2 files changed, 4 insertions(+)
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();
}