My neural network architecture is as follows:
CONV1 : 3 inputs, 64 outputs
POOL1 : 2x2 maxpool 
CONV2 : 64 inputs, 128 outputs
POOL2: 2x2 maxpool 
CONV3: 128 inputs, 256 outputs
CONV4: 256 inputs, 256 outputs
POOL3: 2x2 maxpool 
FC1: 256 inputs, 1024 outputs
FC2: 1024 inputs, 1024 outputs
Batch norm: input= FC2 
SOFTMAX: 1024 inputs, 10 outputs!

My network code:
code|658x340
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
        self.conv3 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1)
        self.conv4 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
        self.fc1   = nn.Linear(256,1024)
        self.fc2   = nn.Linear(1024,1024)
        self.fc2_bn= nn.BatchNorm2d(1024)
        self.classifier = nn.Linear(1024, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), kernel_size=2, stride=2))
        x = F.relu(F.max_pool2d(self.conv2(x),kernel_size=2, stride=2))
        x = F.relu(self.conv3(x))
        x = F.relu(F.max_pool2d(self.conv4(x), kernel_size=2, stride=2))
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2_bn(self.fc2(x)))
        x = F.relu(self.classifier(x))
        return x

I think the problem is with x = x.view(x.size(0), -1). Can someone help me solve this problem ? What is the correct code ?