Hi
I am performing simple image classification task my net is provided below but whenever I am runiing th code I ma getting the follwoing error. Can anyone please help me out hot to fix it.
class CNN(torch.nn.Module):
def __init__(self):
    super(CNN, self).__init__()
    # L1 ImgIn shape=(?, 28, 28, 1)
    #    Conv     -> (?, 28, 28, 32)
    #    Pool     -> (?, 14, 14, 32)
    self.layer1 = torch.nn.Sequential(
        torch.nn.Conv2d(3, 12, kernel_size=3, stride=1, padding=1),
        torch.nn.ReLU(),
        torch.nn.MaxPool2d(kernel_size=2, stride=2),
        torch.nn.Dropout(p=1 - keep_prob))
    # L2 ImgIn shape=(?, 14, 14, 32)
    #    Conv      ->(?, 14, 14, 64)
    #    Pool      ->(?, 7, 7, 64)
    self.layer2 = torch.nn.Sequential(
        torch.nn.Conv2d(12, 24, kernel_size=3, stride=1, padding=1),
        torch.nn.ReLU(),
        torch.nn.MaxPool2d(kernel_size=2, stride=2),
        torch.nn.Dropout(p=1 - keep_prob))
    # L3 ImgIn shape=(?, 7, 7, 64)
    #    Conv      ->(?, 7, 7, 128)
    #    Pool      ->(?, 4, 4, 128)
    self.layer3 = torch.nn.Sequential(
        torch.nn.Conv2d(24, 48, kernel_size=3, stride=1, padding=1),
        torch.nn.ReLU(),
        torch.nn.MaxPool2d(kernel_size=2, stride=2, padding=1),
        torch.nn.Dropout(p=1 - keep_prob))

    # L4 FC 4x4x128 inputs -> 625 outputs
    self.fc1 = torch.nn.Linear(4 * 4 * 128, 32, bias=True)
    torch.nn.init.xavier_uniform(self.fc1.weight)
    self.layer4 = torch.nn.Sequential(
        self.fc1,
        torch.nn.ReLU(),
        torch.nn.Dropout(p=1 - keep_prob))
    # L5 Final FC 625 inputs -> 10 outputs
    self.fc2 = torch.nn.Linear(32, 2, bias=True)
    torch.nn.init.xavier_uniform_(self.fc2.weight) # initialize parameters

def forward(self, x):
    out = self.layer1(x)
    out = self.layer2(out)
    out = self.layer3(out)
    out = out.view(out.size(0), -1)   # Flatten them for FC
    out = self.fc1(out)
    out = self.fc2(out)
    return out

instantiate CNN model
model = CNN()
model
learning_rate = 0.001
criterion = torch.nn.CrossEntropyLoss()    # Softmax is internally computed.
optimizer = torch.optim.Adam(params=model.parameters(), lr=learning_rate)
print(‘Training the Deep Learning network …’)
train_cost = []
train_accu = []
training_epochs = 50
total_batch = train_data_size // bs
print(‘Size of the training dataset is {}’.format(train_data_size))
print(‘Size of the testing dataset is {}’.format(test_data_size))
print(‘Batch size is : {}’.format(bs))
print(‘Total number of batches is : {0:2.0f}’.format(total_batch))
print(’\nTotal number of epochs is : {0:2.0f}’.format(training_epochs))
for epoch in range(training_epochs):
avg_cost = 0
for i, (batch_X, batch_Y) in enumerate(train_data_loader):
X = Variable(batch_X)    # image is already size of (28x28), no reshape
Y = Variable(batch_Y)    # label is not one-hot encoded
    optimizer.zero_grad() # <= initialization of the gradients
    
    # forward propagation
    hypothesis = model(X)
    cost = criterion(hypothesis, Y) # <= compute the loss function
    
    # Backward propagation
    cost.backward() # <= compute the gradient of the loss/cost function     
    optimizer.step() # <= Update the gradients
         
    # Print some performance to monitor the training
    prediction = hypothesis.data.max(dim=1)[1]
    train_accu.append(((prediction.data == Y.data).float().mean()).item())
    train_cost.append(cost.item())   
    if i % 200 == 0:
        print("Epoch= {},\t batch = {},\t cost = {:2.4f},\t accuracy = {}".format(epoch+1, i, train_cost[-1], train_accu[-1]))
   
    avg_cost += cost.data / total_batch

print("[Epoch: {:>4}], averaged cost = {:>.9}".format(epoch + 1, avg_cost.item()))

print(‘Learning Finished!’)
_, argmax = torch.max(outputs, 1)
accuracy = (labels == argmax.squeeze()).float().mean()
RuntimeError: size mismatch, m1: [32 x 3888], m2: [2048 x 32] at /pytorch/aten/src/TH/generic/THTensorMath.cpp:197