RuntimeError: mat1 and mat2 shapes cannot be multiplied (64x16 and 1024x64)

vision
1

import torchvision
from torch.utils.data import DataLoader
train_data = torchvision.datasets.CIFAR10(root="…/data", train=True, transform=torchvision.transforms.ToTensor(),
download=True)
test_data = torchvision.datasets.CIFAR10(root="…/data", train=False, transform=torchvision.transforms.ToTensor(),
download=True)

train_data_size = len(train_data)
test_data_size = len(test_data)

train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)

class Tudui(nn.Module):
def init(self):
super(Tudui, self).init()
self.model = nn.Sequential(
nn.Conv2d(3, 32, 5, 1, 2),
nn.MaxPool2d(2),
nn.Conv2d(32, 32, 5, 1, 2),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 5, 1, 2),
nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(6444, 64),
nn.Linear(64, 10)
)

def forward(self, x):
    x = self.model(x)
    return x

tudui = Tudui()

loss_fn = nn.CrossEntropyLoss()
learning_rate = 1e-2
optimizer = torch.optim.SGD(tudui.parameters(), lr=learning_rate)

total_train_step = 0

total_test_step = 0

epoch = 10

for i in range(epoch):
print(f"第{format(i+1)}轮训练开始")

for data in test_data:
    imgs, targets = data
    output = tudui(imgs)
    loss = loss_fn(output, targets)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    total_train_step = total_train_step + 1
    print(f"训练次数format{(total_train_step)},loss:{format(loss)}")
2

The error is most likely raised in nn.Linear(64*4*4, 64) as the input activation seems to use 64 features instead of the expected 1024.
Change the in_features to 64 and it should work.

PS: you can post code snippets by wrapping them into three backticks ```, which makes debugging easier.