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:

```
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 ?