RuntimeError: The size of tensor a (197) must match the size of tensor b (577) at non-singleton dimension 1

[jeremyEudaric]

Hi,
Congratulation for your amazing work ! when l am using the VIT pre trained model l got this weird issue. Do you have an idea where this issue comes from ?

RuntimeError: The size of tensor a (197) must match the size of tensor b (577) at non-singleton dimension 1

/tmp/ipykernel_701421/2011775653.py in training(epochs)
34
35
---> 36 output = Model(inputs)
37
38 loss = MSE(output,label)

[skynbe]

set ViT('B_16_imagenet1k', pretrained=True, image_size=224)