check attribute existence before access it. attribute is not always traced in our case.
PiperOrigin-RevId: 301904777
Change-Id: Iaf15ad3443354b7f183ffd34f678827d8536ae3a
diff --git a/tensorflow/core/grappler/costs/op_level_cost_estimator.cc b/tensorflow/core/grappler/costs/op_level_cost_estimator.cc
index fdcb892..b427da6 100644
--- a/tensorflow/core/grappler/costs/op_level_cost_estimator.cc
+++ b/tensorflow/core/grappler/costs/op_level_cost_estimator.cc
@@ -1399,7 +1399,9 @@
// Then, the operation to estimate is BatchMatMul([B,M,K],[B,K,N])
const auto& op_info = op_context.op_info;
- string equation = op_info.attr().at("equation").s();
+ auto it = op_info.attr().find("equation");
+ if (it == op_info.attr().end()) return Costs::ZeroCosts(/*inaccurate=*/true);
+ const string& equation = it->second.s();
std::vector<string> equation_split = absl::StrSplit(equation, "->");
if (equation_split.empty()) {
