visualize: fix node_attr on upstream master (#301)

Ref: https://github.com/dmlc/nnvm/pull/152/files#diff-8fb4ac5650c04b74642d321b992f8db0R82
dmlc · Oct 31, 2017 · 41f8f55 · 41f8f55
1 parent 3824df0
commit 41f8f55
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/src/visualize.jl b/src/visualize.jl
@@ -51,10 +51,11 @@ function to_graphviz(network :: SymbolicNode; title="Network Visualization", inp
     attr  = deepcopy(node_attr)
     label = op
 
-    # Up to 0.8 version of mxnet additional info was stored in
-    # node["param"]. Staring from pre0.9 `param` was changed to `attr`.
-    if haskey(node, "param")
-      node_info = node["param"]
+    # Up to 0.11.0 version of mxnet additional info was stored in
+    # node["attr"]. Staring from 0.12 `attr` was changed to `attrs`.
+    # See: https://github.com/dmlc/nnvm/pull/152
+    if haskey(node, "attrs")
+      node_info = node["attrs"]
     elseif haskey(node, "attr")
       node_info = node["attr"]
     end

diff --git a/test/unittest/visualize.jl b/test/unittest/visualize.jl
@@ -12,7 +12,7 @@ function test_basic()
   info("Visualize::basic")
 
   mlp = mlp2()
-  
+
   # Order of elements or default color values can change, but length of the output should be more or less stable
   @test length(mx.to_graphviz(mlp)) == length(
 """