fix : acc test

image_ref/main.py

@@ -58,7 +58,7 @@ def test_duo(model, data_test, loss_function, epoch):
             img_ref = img_ref.cuda()
             label = label.cuda()
         pred_logits = model.forward(imaer,imana,img_ref)
-        pred_class = torch.argmax(pred_logits[:,0],dim=1)
+        pred_class = torch.argmax(pred_logits[:,0])
         acc += (pred_class==label).sum().item()
         loss = loss_function(pred_logits,label)
         losses += loss.item()
@@ -156,8 +156,8 @@ def make_prediction_duo(model, data, f_name, f_name2):
             label = label.cuda()
         output = model(imaer,imana,img_ref)
         confidence_pred_list[specie].append(output[:,0].data.cpu().numpy())
-        #Mono class output (only most postive paire
-        output = torch.argmax(output[:,0],dim=1).data.cpu().numpy()
+        #Mono class output (only most postive paire)
+        output = torch.argmax(output[:,0]).data.cpu().numpy()
         y_pred.extend(output)
 
         label = torch.argmin(label).data.cpu().numpy()