From cf56fefcee35f31c9d69fab5fe8dcea25c46f4d5 Mon Sep 17 00:00:00 2001
From: Franck Dary <franck.dary@lis-lab.fr>
Date: Tue, 23 Jun 2020 16:16:08 +0200
Subject: [PATCH] Corrected bug on CustomHingeLoss where it would not run on
gpu
---
torch_modules/src/CustomHingeLoss.cpp | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/torch_modules/src/CustomHingeLoss.cpp b/torch_modules/src/CustomHingeLoss.cpp
index bc3fc85..ec44b6c 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++)
{
--
GitLab