From 19200d4ec5b058e8c2f6b09204237c42b09acb14 Mon Sep 17 00:00:00 2001
From: Yuanzhong Xu <yuanzx@google.com>
Date: Wed, 14 Oct 2020 23:01:29 -0700
Subject: [PATCH] Fix shape inference call in compile op
PiperOrigin-RevId: 337246090
Change-Id: I7053df6c4d7bb61244ad0b0081904156bffa975a
---
tensorflow/core/tpu/kernels/tpu_compile_op_common.cc | 3 +--
1 file changed, 1 insertion(+), 2 deletions(-)
diff --git a/tensorflow/core/tpu/kernels/tpu_compile_op_common.cc b/tensorflow/core/tpu/kernels/tpu_compile_op_common.cc
index 16f960acefd..eeb396349bb 100644
--- a/tensorflow/core/tpu/kernels/tpu_compile_op_common.cc
+++ b/tensorflow/core/tpu/kernels/tpu_compile_op_common.cc
@@ -514,8 +514,7 @@ Status TpuCompileOpKernelCommon::OptimizeGraph(
// Converts the GraphShapeInfo into the form needed by the constant-folding
// pass of the optimizer.
std::unordered_map<string, std::vector<PartialTensorShape>> shape_map;
- TF_RETURN_IF_ERROR(RunShapeInferenceOnComputation(
- metadata, arg_shapes, graph->get(), flr, &shape_info));
+ ConvertGraphShapeInfoToShapeMap(**graph, shape_info, &shape_map);
optimizer_opts.shape_map = &shape_map;
optimizer.Optimize(flr, flr->env(), flr->device(), graph, optimizer_opts);
}