Make summarize_graph recognize VariableV2 (tensorflow#8201)

pxoo · Mar 15, 2017 · d1aea21 · d1aea21
1 parent d9c61b8
commit d1aea21
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/tensorflow/tools/graph_transforms/summarize_graph_main.cc b/tensorflow/tools/graph_transforms/summarize_graph_main.cc
@@ -109,7 +109,7 @@ Status SummarizeGraph(const GraphDef& graph, const string& graph_path) {
     if (node.op() == "Placeholder") {
       placeholders.push_back(&node);
     }
-    if (node.op() == "Variable") {
+    if (node.op() == "Variable" || node.op() == "VariableV2") {
       variables.push_back(&node);
     }
   }
@@ -168,7 +168,8 @@ Status SummarizeGraph(const GraphDef& graph, const string& graph_path) {
     if (node.device() != "") {
       ++device_counts[node.device()];
     }
-    if ((node.op() == "Const") || (node.op() == "Variable")) {
+    if ((node.op() == "Const") || (node.op() == "Variable") ||
+        (node.op() == "VariableV2")) {
       Tensor tensor;
       if (node.attr().count("value") &&
           tensor.FromProto(node.attr().at("value").tensor())) {