diff --git a/image_ref/main.py b/image_ref/main.py
index dc6b1f6099b833ccf42bf35a1bc4ae24549f5f04..80ee5e1a1d1071e8debd2fd58580d72926ba870d 100644
--- a/image_ref/main.py
+++ b/image_ref/main.py
@@ -153,18 +153,21 @@ def make_prediction_duo(model, data, f_name, f_name2):
         label = label.long()
         specie = torch.argmin(label)
 
+
         if torch.cuda.is_available():
             imaer = imaer.cuda()
             imana = imana.cuda()
             img_ref = img_ref.cuda()
             label = label.cuda()
         output = model(imaer,imana,img_ref)
+        print(output, label)
         confidence_pred_list[specie].append(output[:,0].data.cpu().numpy())
         #Mono class output (only most postive paire)
-        output = torch.argmax(output[:,0]).data.cpu().numpy()
+        output = torch.argmax(output[:,0])
+
         y_pred.extend(output)
 
-        label = torch.argmin(label).data.cpu().numpy()
+        label = torch.argmin(label)
         y_true.extend(label)  # Save Truth
     # constant for classes