diff --git a/image_ref/main.py b/image_ref/main.py
index 003a807658a9533054042d6f18643af92380e214..062e5385eb891b57a12928f0af07160cee26a403 100644
--- a/image_ref/main.py
+++ b/image_ref/main.py
@@ -47,10 +47,11 @@ def test_duo(model, data_test, loss_function, epoch):
         param.requires_grad = False
 
     for imaer,imana, img_ref, label in data_test:
-        imaer=imaer.transpose(0,1)
-        imana=imana.transpose(0,1)
-        img_ref=img_ref.transpose(0,1)
-        label=label.transpose(0,1)
+        imaer = imaer.transpose(0,1)
+        imana = imana.transpose(0,1)
+        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():