I'm struggling with pytorch error

Hi, I am new to pytorch and trying to make CNN model.
I got a help by chatGPT, so I didn’t sincerely understand the whole code.

I got the ‘Expected input batch_size () to match target batch_size ()’ error until yesterday, and now I get the ‘mat1 and mat2 shapes cannot be multiplied (50176x128 and 8192x512)’ error.

I didn’t change the code today, so I don’t know what make this problem.
Also, the batch size error hadn’t been solved for a week. I tried all of the solution I had found, but they didn’t work at all.

Below is my whole code.

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1)
        self.conv3 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=128*8*8, out_features=512)
        self.fc2 = nn.Linear(in_features=512, out_features=10)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = self.pool(x)
        x = self.conv2(x)
        x = self.relu(x)
        x = self.pool(x)
        x = self.conv3(x)
        x = self.relu(x)
        x = self.pool(x)
        x = x.view(-1, 128)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

batch_size = 64

num_workers = 2

train_size = batch_size * (len(dataset) // batch_size)

test_size = len(dataset) - train_size

train_dataset, test_dataset = random_split(dataset, [train_size, test_size])

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, pin_memory=True, num_workers=num_workers)

test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, pin_memory=True, num_workers=num_workers)

model = CNN()

criterion = nn.CrossEntropyLoss()

optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)

lr = 0.001

momentum = 0.9

num_epochs = 10

device = 'cuda' if torch.cuda.is_available() else 'cpu'

model.to(device)

for epoch in range(num_epochs):
    model.train()
    train_loss = 0.0
    train_correct = 0
    total = 0
    for batch_idx, (inputs, targets) in enumerate(train_loader):
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()

        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        train_loss += loss.item() * targets.size(0)
        _, predicted = outputs.max(1)
        total += targets.size(0)
        train_correct += predicted.eq(targets).sum().item()

    train_loss = train_loss / train_size
    train_acc = train_correct / total

    model.eval()
    test_loss = 0.0
    test_correct = 0
    total = 0
    with torch.no_grad():
        for batch_idx, (inputs, targets) in enumerate(test_loader):
            inputs, targets = inputs.to(device), targets.to(device)

            outputs = model(inputs)
            loss = criterion(outputs, targets)

            test_loss += loss.item() * targets.size(0)
            _, predicted = outputs.max(1)
            total += targets.size(0)
            test_correct += predicted.eq(targets).sum().item()

    test_loss = test_loss / test_size
    test_acc = test_correct / total

    print('Epoch [{}/{}], Train Loss: {:.4f}, Train Acc: {:.4f}, Test Loss: {:.4f}, Test Acc: {:.4f}'.format(
        epoch+1, num_epochs, train_loss, train_acc, test_loss, test_acc))

Replace x = x.view(-1, 128) with x = x.view(x.size(0), -1) and fix potential shape mismatch errors in self.fc1 by adapting its in_features to the reported value in the error message.