diff --git a/image_ref/main.py b/image_ref/main.py
index acb26e5592a52c6aa76d5bd603d74290c466583f..8b6b09574084ba067e06318d89df9f87f07fba6c 100644
--- a/image_ref/main.py
+++ b/image_ref/main.py
@@ -20,7 +20,7 @@ def train_duo(model, data_train, optimizer, loss_function, epoch):
 
     for imaer,imana, img_ref, label in data_train:
         label = label.long()
-        print(imaer.size)
+        print(imaer.shape)
         if torch.cuda.is_available():
             imaer = imaer.cuda()
             imana = imana.cuda()
@@ -51,7 +51,7 @@ def test_duo(model, data_test, loss_function, epoch):
         imana.transpose(0,1)
         img_ref.transpose(0,1)
         label.transpose(0,1)
-        print(imaer.size)
+        print(imaer.shape)
         label = label.long()
         if torch.cuda.is_available():
             imaer = imaer.cuda()