code cleanup

demo_graph.py

@@ -25,7 +25,7 @@ def forward(self, x):
         x = F.relu(self.fc1(x))
         x = F.dropout(x, training=self.training)
         x = self.fc2(x)
-        x = F.softmax(x)
+        x = F.softmax(x, dim=1)
         return x
 
 
@@ -46,7 +46,7 @@ def forward(self, x):
         x = F.dropout(x, training=self.training)
         x = self.fc2(x)
         # x = F.log_softmax(x) no scopename(JIT bug)
-        x = F.softmax(x)
+        x = F.softmax(x, dim=1)
         return x
 
 dummy_input = Variable(torch.rand(13, 1, 28, 28))

tensorboardX/graph.py

@@ -3,11 +3,9 @@
 from .src.versions_pb2 import VersionDef
 from .src.attr_value_pb2 import AttrValue
 from .src.tensor_shape_pb2 import TensorShapeProto
-import torch
 
 
 def replace(name, scope):
-    print(type(name), name)
     return '/'.join([scope[name], name])
 
 
@@ -32,24 +30,22 @@ def parse(graph):
         nodes.append({'name': replace(uname, scope), 'op': n.kind(), 'inputs': inputs, 'attr': attrs})
 
     for n in graph.inputs():
-        print(n.type())
         uname = n.uniqueName()
         nodes.append({'name': replace(uname, scope), 'op': 'Parameter', 'inputs': [], 'attr': str(n.type())})
 
     return nodes
 
 
 def graph(model, args):
+    import torch
     with torch.onnx.set_training(model, False):
         trace, _ = torch.jit.trace(model, args)
     torch.onnx._optimize_trace(trace, False)
     graph = trace.graph()
-    print(graph)
     list_of_nodes = parse(graph)
     nodes = []
     for node in list_of_nodes:
         nodes.append(
             NodeDef(name=node['name'], op=node['op'], input=node['inputs'],
                     attr={'lanpa': AttrValue(s=node['attr'].encode(encoding='utf_8'))}))
-    print(nodes)
     return GraphDef(node=nodes, versions=VersionDef(producer=22))