Hi. I’m new to deep learing. I, according to an example adapted to the MNIST,  try something new. I import kaggle kernel Images of Lego Bricks 2. But there is a different size of the pictures. And I get an error.

shape ‘[-1, 256]’ is invalid for input of size 141376

I changed 16 * 4 *4 to 16 * 94 *94, but got new error

Expected input batch_size (1) to match target batch_size (4).

I changed batch_size to 1 and got new one

shape ‘[-1, 141376]’ is invalid for input of size 35344

It seems to me that the neural network is not correct. Please tell me how to design a neural network for this case. Thanks.
P.S. If you have an article explaining how to do this, give a link to it.
import torch
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import torch.optim as optim
import numpy as np
import os
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 4 * 4, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 4 *4)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

classes = os.listdir("F:\\projects\\conda\\datasets\\lego-brick-images\\train")
batch_size = 4
epochs=30
lr=0.01
log_interval=1000

train_dataset = torchvision.datasets.ImageFolder(root="F:\\projects\\conda\\datasets\\lego-brick-images\\train", 
                                                 transform=transforms.ToTensor())
test_dataset = torchvision.datasets.ImageFolder(root="F:\\projects\\conda\\datasets\\lego-brick-images\\valid", 
                                                transform=transforms.ToTensor())

trainloader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size,
                                          shuffle=True, num_workers=2)
testloader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size,
                                         shuffle=False, num_workers=2)

model = Net();
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9)  

def train(model, train_loader, optimizer, epoch, log_interval):
    model.train()
    avg_loss = 0
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad() 
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        avg_loss+=F.nll_loss(output, target, reduction='sum').item()
        
        if batch_idx % log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))
    avg_loss/=len(train_loader.dataset)
    return avg_loss

def test(model, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            output = model(data)
            test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))
    accuracy = 100. * correct / len(test_loader.dataset)
    return test_loss,accuracy

train_losses = []
test_losses = []
accuracy_list = []
for epoch in range(1, epochs + 1):
    trn_loss = train(model, trainloader, optimizer, epoch, log_interval)
    test_loss, accuracy = test(model, testloader)
    train_losses.append(trn_loss)
    test_losses.append(test_loss)
    accuracy_list.append(accuracy)