From fbc490b2f02d2274e6864541f86053d12884d0b9 Mon Sep 17 00:00:00 2001
From: Lluis-Miquel Munguia <llmunguia@google.com>
Date: Thu, 22 Aug 2019 14:36:57 -0700
Subject: [PATCH] Add memory usage details to the cost graph.
PiperOrigin-RevId: 264921808
---
tensorflow/core/grappler/costs/analytical_cost_estimator.cc | 2 ++
1 file changed, 2 insertions(+)
diff --git a/tensorflow/core/grappler/costs/analytical_cost_estimator.cc b/tensorflow/core/grappler/costs/analytical_cost_estimator.cc
index a85e293ac00..5a79441eee9 100644
--- a/tensorflow/core/grappler/costs/analytical_cost_estimator.cc
+++ b/tensorflow/core/grappler/costs/analytical_cost_estimator.cc
@@ -61,6 +61,8 @@ void AddCostNode(ReadyNodeManager* node_manager, const OpContext& op_context,
node->set_compute_cost(node_costs.execution_time.asMicroSeconds().count());
node->set_compute_time(node_costs.compute_time.asMicroSeconds().count());
node->set_memory_time(node_costs.memory_time.asMicroSeconds().count());
+ node->set_temporary_memory_size(node_costs.temporary_memory);
+ node->set_persistent_memory_size(node_costs.persistent_memory);
node->set_inaccurate(node_costs.inaccurate);
for (const string& input : node_manager->GetCurrNode()->input()) {