Merge branch 'master' into jbobba/batchnorm-layout

d7bd9bb1 · Jayaram Bobba · GitHub · 1c27b622 · dd3af46c · d7bd9bb1
Unverified Commit d7bd9bb1 authored Feb 23, 2018 by Jayaram Bobba Committed by GitHub Feb 23, 2018
Hide whitespace changes
Inline Side-by-side

Showing with 10 additions and 2 deletions

visualize_tree.cpp src/ngraph/pass/visualize_tree.cpp +10 -2

No files found.
--- a/src/ngraph/pass/visualize_tree.cpp
+++ b/src/ngraph/pass/visualize_tree.cpp
@@ -66,9 +66,17 @@ std::string pass::VisualizeTree::add_attributes(shared_ptr<Node> node)
 std::string pass::VisualizeTree::get_attributes(shared_ptr<Node> node)
 {
    stringstream ss;
-    if (node->is_parameter())
+    if (node->is_parameter() || node->is_output())
    {
-        ss << "    " << node->get_name() << " [shape=box color=blue";
+        ss << "    " << node->get_name() << " [shape=box ";
+        if (node->is_parameter())
+        {
+            ss << "color=blue ";
+        }
+        if (node->is_output())
+        {
+            ss << "style=filled fillcolor=pink ";
+        }
    }
    else
    {