Sending torch object to CPU before saving them to disk

reading_machine/src/Classifier.cpp

@@ -229,12 +229,16 @@ std::string Classifier::getLastFilename() const
 
 void Classifier::saveBest()
 {
+  getNN()->to(torch::kCPU);
   torch::save(getNN(), getBestFilename());
+  getNN()->to(NeuralNetworkImpl::device);
 }
 
 void Classifier::saveLast()
 {
+  getNN()->to(torch::kCPU);
   torch::save(getNN(), getLastFilename());
+  getNN()->to(NeuralNetworkImpl::device);
   saveOptimizer();
 }
 

trainer/src/Trainer.cpp

@@ -296,7 +296,7 @@ void Trainer::Examples::saveIfNeeded(const std::string & state, std::filesystem:
 
   int nbClasses = classes[0].size(0);
 
-  auto tensorToSave = torch::cat({torch::stack(contexts), torch::stack(classes)}, 1);
+  auto tensorToSave = torch::cat({torch::stack(contexts), torch::stack(classes)}, 1).to(torch::kCPU);
   auto filename = fmt::format("{}-{}_{}-{}.{}.{}.tensor", state, nbClasses, lastSavedIndex, currentExampleIndex-1, epoch, dynamicOracle);
   torch::save(tensorToSave, dir/filename);
   lastSavedIndex = currentExampleIndex;