From ac85f8136b287f231ea495ecba5a72341018a1ae Mon Sep 17 00:00:00 2001
From: Jeff Donahue <jeff.donahue@gmail.com>
Date: Thu, 21 Aug 2014 12:53:43 -0700
Subject: [PATCH] If specified, --gpu flag overrides SolverParameter
 solver_mode.

---
 tools/caffe.cpp | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/tools/caffe.cpp b/tools/caffe.cpp
index 13398ae..eb32128 100644
--- a/tools/caffe.cpp
+++ b/tools/caffe.cpp
@@ -89,6 +89,17 @@ int train() {
 
   LOG(INFO) << "Starting Optimization";
   caffe::SGDSolver<float> solver(solver_param);
+
+  // Set device id and mode
+  if (FLAGS_gpu >= 0) {
+    LOG(INFO) << "Use GPU with device ID " << FLAGS_gpu;
+    Caffe::SetDevice(FLAGS_gpu);
+    Caffe::set_mode(Caffe::GPU);
+  } else if (!solver_param.has_solver_mode()) {
+    LOG(INFO) << "Use CPU.";
+    Caffe::set_mode(Caffe::CPU);
+  }
+
   if (FLAGS_snapshot.size()) {
     LOG(INFO) << "Resuming from " << FLAGS_snapshot;
     solver.Solve(FLAGS_snapshot);
-- 
2.7.4