Can't figure out the error?

username: ZainUlHaq
cooked: <p>class ImageClassification(nn.Module):</p>
<pre><code>def __init__(self):

    super().__init__()

    self.conv1 = nn.Conv3d(1, 8, 3, 2) 

    self.conv1_bn = nn.BatchNorm3d(8)

    self.conv2 = nn.Conv3d(8, 16, 3, 1)

    self.conv2_bn = nn.BatchNorm3d(16)

    self.conv3 = nn.Conv3d(16,32,3,1)

    self.conv3_bn = nn.BatchNorm3d(32)

    self.fc1 = nn.Linear(4*4*1*32, 256)

    self.fc1_bn = nn.BatchNorm3d(256)

    self.fc2 = nn.Linear(256, 128) 

    self.fc2_bn = nn.BatchNorm3d(128)

    self.fc3 = nn.Linear(128, 2)

def forward(self, X):

    X = self.conv1(X)

    X = F.relu(self.conv1_bn(X))

    X = F.max_pool3d(X,2)

    X = self.conv2(X)

    X = F.relu(self.conv2_bn(X))

    X = F.max_pool3d(X,2)

    X = self.conv3(X)

    X = F.relu(self.conv3_bn(X))

    X = F.max_pool3d(X,2)

    X = x.view(-1,512)

    X = self.fc1(X)

    X = self.fc1_bn(X)

    X = self.fc2(X)

    X = F.relu(self.fc2_bn(X))

    X=  F.relu(self.fc3(X))

    return F.log_softmax(X, dim=1)
</code></pre>
<hr>
<p>ValueError                                Traceback (most recent call last)<br>
 in ()<br>
17<br>
18         # Apply the model<br>
—&gt; 19         y_pred = model(X_train)  # we don’t flatten X-train here<br>
20         loss = criterion(y_pred, y_train)<br>
21</p>
<p>4 frames<br>
/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)<br>
1100         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks<br>
1101                 or _global_forward_hooks or _global_forward_pre_hooks):<br>
 → 1102             return forward_call(*input, **kwargs)<br>
1103         # Do not call functions when jit is used<br>
1104         full_backward_hooks, non_full_backward_hooks = [], []</p>
<p> in forward(self, X)<br>
26         X = x.view(-1,512)<br>
27         X = self.fc1(X)<br>
—&gt; 28         X = self.fc1_bn(X)<br>
29         X = self.fc2(X)<br>
30         X = F.relu(self.fc2_bn(X))</p>
<p>/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)<br>
1100         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks<br>
1101                 or _global_forward_hooks or _global_forward_pre_hooks):<br>
 → 1102             return forward_call(*input, **kwargs)<br>
1103         # Do not call functions when jit is used<br>
1104         full_backward_hooks, non_full_backward_hooks = [], []</p>
<p>/usr/local/lib/python3.7/dist-packages/torch/nn/modules/batchnorm.py in forward(self, input)<br>
133<br>
134     def forward(self, input: Tensor) → Tensor:<br>
 → 135         self._check_input_dim(input)<br>
136<br>
137         # exponential_average_factor is set to self.momentum</p>
<p>/usr/local/lib/python3.7/dist-packages/torch/nn/modules/batchnorm.py in _check_input_dim(self, input)<br>
512     def _check_input_dim(self, input):<br>
513         if input.dim() != 5:<br>
 → 514             raise ValueError(“expected 5D input (got {}D input)”.format(input.dim()))<br>
515<br>
516</p>
<p>ValueError: expected 5D input (got 2D input)</p>

username: thecho7
cooked: <p>Let me explain what did you wrong in 2 steps,</p>
<blockquote>
<p>X = x.view(-1, 512)<br>
X = self.fc1(X)</p>
</blockquote>
<p>X from the above has a size of (batch_size, 256) as you define the output size of self.fc1 256.</p>
<p>Next,</p>
<blockquote>
<p>X = self.fc1_bn(X)</p>
</blockquote>
<p>The input X is a 2-dimensional tensor but fc1_bn is defined as 3-D batch normalization layer.<br>
If you really want to use 3-D BatchNorm, then you should reform X before the normalization e.g.</p>
<blockquote>
<p>X = X.view(1, -1, 1, 1, 256)<br>
X = self.fc1_bn(X)</p>
</blockquote>
<p>I recommend you not use batch normalization after nn.Linear. It means nothing with 2-dimensional tensor.<br>
Just an activation is enough such as ReLU, Sigmoid, and so on.<br>
Which activation function is the best? It’s totally up to your task.</p>