Skip to content

Commit

Permalink
python: Set gpu device id before setting gpu mode
Browse files Browse the repository at this point in the history
Otherwise caffe allocates some memory on GPU#0
  • Loading branch information
xerus committed Aug 14, 2018
1 parent 2a1c552 commit 4353628
Show file tree
Hide file tree
Showing 2 changed files with 2 additions and 2 deletions.
2 changes: 1 addition & 1 deletion docs/tutorial/interfaces.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion python/train.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down

0 comments on commit 4353628

Please sign in to comment.