Merge pull request BVLC#521 from sguada/set_device_id_at_init

Set device_id at the beginning of Solver.Init()
mitmul · Jun 20, 2014 · b1e1e23 · b1e1e23
2 parents 1726257 + 29f593c
commit b1e1e23
Showing 1 changed file with 5 additions and 6 deletions.
diff --git a/src/caffe/solver.cpp b/src/caffe/solver.cpp
@@ -36,6 +36,11 @@ void Solver<Dtype>::Init(const SolverParameter& param) {
   LOG(INFO) << "Initializing solver from parameters: " << std::endl
             << param.DebugString();
   param_ = param;
+  if (param_.solver_mode() == SolverParameter_SolverMode_GPU &&
+      param_.has_device_id()) {
+    Caffe::SetDevice(param_.device_id());
+  }
+  Caffe::set_mode(Caffe::Brew(param_.solver_mode()));
   if (param_.random_seed() >= 0) {
     Caffe::set_random_seed(param_.random_seed());
   }
@@ -74,14 +79,8 @@ void Solver<Dtype>::Init(const SolverParameter& param) {
   LOG(INFO) << "Solver scaffolding done.";
 }
 
-
 template <typename Dtype>
 void Solver<Dtype>::Solve(const char* resume_file) {
-  Caffe::set_mode(Caffe::Brew(param_.solver_mode()));
-  if (param_.solver_mode() == SolverParameter_SolverMode_GPU &&
-      param_.has_device_id()) {
-    Caffe::SetDevice(param_.device_id());
-  }
   Caffe::set_phase(Caffe::TRAIN);
   LOG(INFO) << "Solving " << net_->name();
   PreSolve();