Ошибка size mismatch, m1: [7168 x 56], m2: [401408 x 56] at /pytorch/aten/src/TH/generic/THTensorMath.cpp:41

Моя модель выглядит следующим образом, я понимаю, что ошибка в слое fc1. На вход у меня подаются изображения 224*224, размер батча 8.

Я посмотрела вывод из предыдущего слоя bn2 и он выглядит так: torch.Size([8, 16, 56, 56]). Поэтому я не понимаю в чем моя ошибка. Помогите, пожалуйста

class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(in_channels=6, out_channels=6, kernel_size=3, padding=1)
        self.act1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.bn1 = torch.nn.BatchNorm2d(num_features=6)
        
        self.conv3 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=3, padding=1)
        self.conv4 = nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, padding=1)
        self.act2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.bn2 = torch.nn.BatchNorm2d(num_features=16)
        
        self.fc1 = nn.Linear(in_features=56 * 56 * 16 * 8, out_features=56)
        self.fc2 = nn.Linear(56, 10)
        

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.act1(x)
        x = self.pool1(x)
        x = self.bn1(x)
        
        x = self.conv3(x)
        x = self.conv4(x)
        x = self.act2(x)
        x = self.pool2(x)
        print(x.shape)
        x = self.bn2(x)
        print(x.shape)
        x = self.fc1(x)
        print(x.shape)
        x = self.fc2(x)
        
        return x

Полная ошибка:

RuntimeError                              Traceback (most recent call last)
<ipython-input-44-a7c73a6166af> in <module>
----> 1 accuracy_train, loss_train = train(15)
      2 #

<ipython-input-43-361b75daf0f2> in train(n_epoch)
     15             optimizer.zero_grad()
     16 
---> 17             preds = network(inputs)
     18             losses = loss(preds, labels)
     19             losses.backward()

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    720             result = self._slow_forward(*input, **kwargs)
    721         else:
--> 722             result = self.forward(*input, **kwargs)
    723         for hook in itertools.chain(
    724                 _global_forward_hooks.values(),

<ipython-input-42-07d9510bd54f> in forward(self, x)
     34         x = self.bn2(x)
     35         print(x.shape)
---> 36         x = self.fc1(x)
     37         print(x.shape)
     38         x = self.fc2(x)

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    720             result = self._slow_forward(*input, **kwargs)
    721         else:
--> 722             result = self.forward(*input, **kwargs)
    723         for hook in itertools.chain(
    724                 _global_forward_hooks.values(),

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/linear.py in forward(self, input)
     89 
     90     def forward(self, input: Tensor) -> Tensor:
---> 91         return F.linear(input, self.weight, self.bias)
     92 
     93     def extra_repr(self) -> str:

/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py in linear(input, weight, bias)
   1674         ret = torch.addmm(bias, input, weight.t())
   1675     else:
-> 1676         output = input.matmul(weight.t())
   1677         if bias is not None:
   1678             output += bias

RuntimeError: size mismatch, m1: [7168 x 56], m2: [401408 x 56] at /pytorch/aten/src/TH/generic/THTensorMath.cpp:41 ```