Workaround for running on default GPU if ID not specified

cypof · Jul 14, 2015 · b060aa1 · b060aa1
1 parent 48a7adf
commit b060aa1
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/tools/caffe.cpp b/tools/caffe.cpp
@@ -134,9 +134,13 @@ int train() {
   // If the gpus flag is not provided, allow the mode and device to be set
   // in the solver prototxt.
   if (FLAGS_gpu.size() == 0
-      && solver_param.solver_mode() == caffe::SolverParameter_SolverMode_GPU
-      && solver_param.has_device_id()) {
-    FLAGS_gpu = "" + boost::lexical_cast<string>(solver_param.device_id());
+      && solver_param.solver_mode() == caffe::SolverParameter_SolverMode_GPU) {
+      if (solver_param.has_device_id()) {
+          FLAGS_gpu = ""  +
+              boost::lexical_cast<string>(solver_param.device_id());
+      } else {  // Set default GPU if unspecified
+          FLAGS_gpu = "" + boost::lexical_cast<string>(0);
+      }
   }
 
   vector<int> gpus;