fix(torch): reload solver params on API device

jolibrain · Nov 5, 2020 · 30fa16f · 30fa16f
1 parent 5ab90c7
commit 30fa16f
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 5 deletions.
diff --git a/src/backends/torch/torchlib.cc b/src/backends/torch/torchlib.cc
@@ -471,7 +471,7 @@ namespace dd
 
     int it = 0;
     // reload solver and set it value accordingly
-    it = tsolver.load(this->_mlmodel._sstate);
+    it = tsolver.load(this->_mlmodel._sstate, _device);
     tsolver.zero_grad();
     _module.train();
 

diff --git a/src/backends/torch/torchmodule.cc b/src/backends/torch/torchmodule.cc
@@ -75,6 +75,7 @@ namespace dd
     try
       {
         _graph = std::make_shared<CaffeToTorch>(model._proto);
+        _graph->to(_device);
       }
     catch (std::exception &e)
       {
@@ -107,7 +108,7 @@ namespace dd
         _logger->info("loading " + tmodel._native);
         try
           {
-            torch::load(_native, tmodel._native);
+            torch::load(_native, tmodel._native, _device);
           }
         catch (std::exception &e)
           {

diff --git a/src/backends/torch/torchsolver.cc b/src/backends/torch/torchsolver.cc
@@ -128,7 +128,7 @@ namespace dd
     torch::save(*_optimizer, sfile);
   }
 
-  int TorchSolver::load(std::string sstate)
+  int TorchSolver::load(std::string sstate, torch::Device device)
   {
     if (!sstate.empty())
       {
@@ -140,7 +140,7 @@ namespace dd
         _logger->info("loading " + sstate);
         try
           {
-            torch::load(*_optimizer, sstate);
+            torch::load(*_optimizer, sstate, device);
           }
         catch (std::exception &e)
           {

diff --git a/src/backends/torch/torchsolver.h b/src/backends/torch/torchsolver.h
@@ -42,7 +42,7 @@ namespace dd
     void configure(APIData ad_solver);
     void create(TorchModule &module);
 
-    int load(std::string sstate);
+    int load(std::string sstate, torch::Device device);
     void save(std::string sfile);
 
     void zero_grad()