fix : dataloader batched

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 = []