一部のパラメータのみに名前をつけられるように変更

from szagoruyko#21 (comment)
juri200405 · Dec 22, 2020 · be8b7dd · be8b7dd
1 parent 46add7f
commit be8b7dd
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchviz/dot.py b/torchviz/dot.py
@@ -44,7 +44,7 @@ def add_nodes(var):
                 dot.node(str(id(var)), size_to_str(var.size()), fillcolor='orange')
             elif hasattr(var, 'variable'):
                 u = var.variable
-                name = param_map[id(u)] if params is not None else ''
+                name = param_map.get(id(u), ' ') if params is not None else ''
                 node_name = '%s\n %s' % (name, size_to_str(u.size()))
                 dot.node(str(id(var)), node_name, fillcolor='lightblue')
             elif var in output_nodes: