I want to use the pretrained model in torchvision to finish the cats VS dog tasks. But there is something wrong with my code : (
RuntimeError: size mismatch, m1: [2 x 512], m2: [25088 x 4096] at /pytorch/aten/src/THC/generic/THCTensorMathBlas.cu:249
The problem seem comes from the net structure. But I just load the model in the torchvision.
Can anyone give me some suggestions?
And, python == 3.6 torch == 0.4.0, torchvision == 0.2.1.
import torch import torch.nn as nn import torchvision.models as models import torchvision.transforms as transforms net = models.vgg16(pretrained=True) net = net.to(device) test = torch.randn(2,3,32,32).to(device) y = net(test)