diff --git a/image_ref/main.py b/image_ref/main.py index dba220ad7967946ee559295f515fd575d3eb79e1..5cac8f3d24da4c33bc62f49e898db342ed063b06 100644 --- a/image_ref/main.py +++ b/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()