fix draw_net python script

include caffe.draw for drawing functions.
marshallixp · May 22, 2014 · 1cb6fea · 1cb6fea
1 parent 2cb1359
commit 1cb6fea
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/draw_net.py b/python/draw_net.py
@@ -5,7 +5,7 @@
 import os
 from google.protobuf import text_format
 
-import caffe
+import caffe, caffe.draw
 from caffe.proto import caffe_pb2
 
 
@@ -17,7 +17,7 @@ def main(argv):
         net = caffe_pb2.NetParameter()
         text_format.Merge(open(sys.argv[1]).read(), net)
         print 'Drawing net to %s' % sys.argv[2]
-        draw_net_to_file(net, sys.argv[2])
+        caffe.draw.draw_net_to_file(net, sys.argv[2])
 
 
 if __name__ == '__main__':