From 1795d057d779802f996bc10a4d61b5724bd17bef Mon Sep 17 00:00:00 2001
From: Schneider Leo <leo.schneider@etu.ec-lyon.fr>
Date: Wed, 12 Mar 2025 16:24:33 +0100
Subject: [PATCH] model duo cuda error fix
---
main.py | 3 +++
1 file changed, 3 insertions(+)
diff --git a/main.py b/main.py
index 84240bf..db56d0f 100644
--- a/main.py
+++ b/main.py
@@ -137,6 +137,7 @@ def train_duo(model, data_train, optimizer, loss_function, epoch):
if torch.cuda.is_available():
imaer = imaer.cuda()
imana = imana.cuda()
+ label = label.cuda()
pred_logits = model.forward(imaer,imana)
pred_class = torch.argmax(pred_logits,dim=1)
acc += (pred_class==label).sum().item()
@@ -162,6 +163,7 @@ def test_duo(model, data_test, loss_function, epoch):
if torch.cuda.is_available():
imaer = imaer.cuda()
imana = imana.cuda()
+ label = label.cuda()
pred_logits = model.forward(imaer,imana)
pred_class = torch.argmax(pred_logits,dim=1)
acc += (pred_class==label).sum().item()
@@ -220,6 +222,7 @@ def make_prediction_duo(model, data, f_name):
if torch.cuda.is_available():
imaer = imaer.cuda()
imana = imana.cuda()
+ label = label.cuda()
output = model(imaer,imana)
output = (torch.max(torch.exp(output), 1)[1]).data.cpu().numpy()
--
GitLab