diff --git a/image_ref/main.py b/image_ref/main.py index 12ebb403bd5520d65795295c4f856dda55a90d50..acb26e5592a52c6aa76d5bd603d74290c466583f 100644 --- a/image_ref/main.py +++ b/image_ref/main.py @@ -20,6 +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) if torch.cuda.is_available(): imaer = imaer.cuda() imana = imana.cuda() @@ -50,6 +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) label = label.long() if torch.cuda.is_available(): imaer = imaer.cuda()