diff --git a/torch_modules/src/CustomHingeLoss.cpp b/torch_modules/src/CustomHingeLoss.cpp
index bc3fc859cafeb611eea5e28559bf23901843f3c4..ec44b6c5bdba5aac28d314a11bfb06175f1b1f7b 100644
--- a/torch_modules/src/CustomHingeLoss.cpp
+++ b/torch_modules/src/CustomHingeLoss.cpp
@@ -1,8 +1,9 @@
 #include "CustomHingeLoss.hpp"
+#include "NeuralNetwork.hpp"
 
 torch::Tensor CustomHingeLoss::operator()(torch::Tensor prediction, torch::Tensor gold)
 {
-  torch::Tensor loss = torch::zeros(1);
+  torch::Tensor loss = torch::zeros(1).to(NeuralNetworkImpl::device);
 
   for (unsigned int i = 0; i < prediction.size(0); i++)
   {