From 80730f889696ad8e02b4c854d708262c22370298 Mon Sep 17 00:00:00 2001
From: Franck Dary <franck.dary@lis-lab.fr>
Date: Thu, 2 Jul 2020 13:41:37 +0200
Subject: [PATCH] Senting neural network to correct device

---
 reading_machine/src/Classifier.cpp | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/reading_machine/src/Classifier.cpp b/reading_machine/src/Classifier.cpp
index 7023030..c685fa2 100644
--- a/reading_machine/src/Classifier.cpp
+++ b/reading_machine/src/Classifier.cpp
@@ -66,6 +66,8 @@ Classifier::Classifier(const std::string & name, std::filesystem::path path, std
   getNN()->loadDicts(path);
   getNN()->registerEmbeddings();
 
+  getNN()->to(NeuralNetworkImpl::device);
+
   if (!train)
     torch::load(getNN(), getBestFilename());
   else if (std::filesystem::exists(getLastFilename()))
-- 
GitLab