Set cfg.GPU_ID in demo.py

puzzledqs · Oct 28, 2015 · 81ee227 · 81ee227
1 parent 19062d5
commit 81ee227
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tools/demo.py b/tools/demo.py
@@ -131,6 +131,7 @@ def parse_args():
     else:
         caffe.set_mode_gpu()
         caffe.set_device(args.gpu_id)
+        cfg.GPU_ID = args.gpu_id
     net = caffe.Net(prototxt, caffemodel, caffe.TEST)
 
     print '\n\nLoaded network {:s}'.format(caffemodel)