python: Set gpu device id before setting gpu mode

Otherwise caffe allocates some memory on GPU#0
BVLC · Aug 14, 2018 · 4353628 · 4353628
1 parent 2a1c552
commit 4353628
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/docs/tutorial/interfaces.md b/docs/tutorial/interfaces.md
@@ -129,8 +129,8 @@ Use CPU:
 
 Use GPU and specify its gpu_id:
 
-    caffe.set_mode_gpu();
     caffe.set_device(gpu_id);
+    caffe.set_mode_gpu();
 
 #### Create a network and access its layers and blobs
 

diff --git a/python/train.py b/python/train.py
@@ -63,8 +63,8 @@ def show_time():
 
 
 def solve(proto, snapshot, gpus, timing, uid, rank):
-    caffe.set_mode_gpu()
     caffe.set_device(gpus[rank])
+    caffe.set_mode_gpu()
     caffe.set_solver_count(len(gpus))
     caffe.set_solver_rank(rank)
     caffe.set_multiprocess(True)