fix : dataloader batched

image_ref/main.py

@@ -19,7 +19,6 @@ def train_duo(model, data_train, optimizer, loss_function, epoch):
         param.requires_grad = True
 
     for imaer,imana, img_ref, label in data_train:
-        print(imaer.size())
         label = label.long()
         if torch.cuda.is_available():
             imaer = imaer.cuda()
@@ -47,7 +46,10 @@ def test_duo(model, data_test, loss_function, epoch):
         param.requires_grad = False
 
     for imaer,imana, img_ref, label in data_test:
-        print(imaer.size())
+        imaer.permute(0,1)
+        imana.permute(0,1)
+        img_ref.permute(0,1)
+        label.permute(0,1)
         label = label.long()
         if torch.cuda.is_available():
             imaer = imaer.cuda()