From 3e8e55eedad9d66f6d02dffe49e25ae97f0336d3 Mon Sep 17 00:00:00 2001
From: Schneider Leo <leo.schneider@etu.ec-lyon.fr>
Date: Tue, 1 Apr 2025 15:56:54 +0200
Subject: [PATCH] fix : dataloader batched
---
image_ref/main.py | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/image_ref/main.py b/image_ref/main.py
index 062e5385..3c48ef2b 100644
--- a/image_ref/main.py
+++ b/image_ref/main.py
@@ -20,7 +20,6 @@ def train_duo(model, data_train, optimizer, loss_function, epoch):
for imaer,imana, img_ref, label in data_train:
label = label.long()
- print(label.shape)
if torch.cuda.is_available():
imaer = imaer.cuda()
imana = imana.cuda()
@@ -52,7 +51,6 @@ def test_duo(model, data_test, loss_function, epoch):
img_ref = img_ref.transpose(0,1)
label = label.transpose(0,1)
label = label.squeeze()
- print(label.shape)
label = label.long()
if torch.cuda.is_available():
imaer = imaer.cuda()
@@ -121,7 +119,9 @@ def run_duo(args):
def make_prediction_duo(model, data, f_name):
- n_class = len(data[0][2])
+ for imaer, imana, img_ref, label in data:
+ n_class = len(label.shape[1])
+ break
confidence_pred_list = [[] for i in range(n_class)]
y_pred = []
y_true = []
--
GitLab