Erro ao realizar o treinamento com a arquitetura VGG11

Não estou conseguindo realizar a etapa de treino e validação da CNN com a arquitetura descrita no exercício. A arquitetura utilizada é descrita abaixo (da mesma maneira dada como resposta do exercício).

# Definindo a rede
net = nn.Sequential(
        ## ConvBlock 1
        nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),      # entrada: (b, 3, 224, 224) e saida: (b, 64, 224, 224)
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=2, stride=2, padding=0),          # entrada: (b, 64, 224, 224) e saida: (b, 64, 112, 112)

        ## ConvBlock 2
        nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),    # entrada: (b, 64, 112, 112) e saida: (b, 128, 112, 112)
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=2, stride=2, padding=0),          # entrada: (b, 128, 112, 112) e saida: (b, 128, 56, 56)

        ## ConvBlock 3
        nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),   # entrada: (b, 128, 56, 56) e saida: (b, 256, 56, 56)
        nn.ReLU(),
        nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),   # entrada: (b, 256, 56, 56) e saida: (b, 256, 56, 56)
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=2, stride=2, padding=0),          # entrada: (b, 256, 56, 56) e saida: (b, 256, 28, 28)

        ## ConvBlock 4
        nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1),   # entrada: (b, 256, 28, 28) e saida: (b, 512, 28, 28)
        nn.ReLU(),
        nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),   # entrada: (b, 512, 28, 28) e saida: (b, 512, 28, 28)
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=2, stride=2, padding=0),          # entrada: (b, 512, 28, 28) e saida: (b, 512, 14, 14)

        ## ConvBlock 4
        nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),   # entrada: (b, 512, 14, 14) e saida: (b, 512, 14, 14)
        nn.ReLU(),
        nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),   # entrada: (b, 512, 14, 14) e saida: (b, 512, 14, 14)
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=2, stride=2, padding=0),          # entrada: (b, 512, 14, 14) e saida: (b, 512, 7, 7)
        nn.Flatten(),  # lineariza formando um vetor               # entrada: (b, 512, 7, 7) e saida: (b, 512*7*7) = (b, 25088)

        ## DenseBlock
        nn.Linear(25088, 4096),                                    # entrada: (b, 25088) e saida: (b, 4096)
        nn.ReLU(),
        nn.Linear(4096, 4096),                                     # entrada: (b, 4096) e saida: (b, 4096)
        nn.ReLU(),
        nn.Linear(4096, 10),                                       # entrada: (b, 4096) e saida: (b, 10)
        nn.Softmax(dim=-1)
        )

Entretanto, ao rodar na etapa de teste e validação, me deparo com o seguinte erro:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-29-3019b93140ee> in <module>()
      3 
      4   # Train
----> 5   train_losses.append(train(train_loader, net, epoch))
      6 
      7   # Validate

<ipython-input-16-ef27e004f856> in train(train_loader, net, epoch)
     17 
     18     # Forward
---> 19     ypred = net(dado)
     20     loss = criterion(ypred, rotulo)
     21     epoch_loss.append(loss.cpu().data)

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/linear.py in forward(self, input)
    112 
    113     def forward(self, input: Tensor) -> Tensor:
--> 114         return F.linear(input, self.weight, self.bias)
    115 
    116     def extra_repr(self) -> str:

RuntimeError: mat1 and mat2 shapes cannot be multiplied (50x512 and 25088x4096)

Pesquisei em fóruns e, até onde entendi, parece que nn.Linear(25088, 4096) esperava um tensor com 25088 (=51277) mas está recebendo 512.

A pergunta é, como eu poderia arrumar a arquitetura dessa CNN ? Seria necessário acrescentar mais uma camada de convolução para diminuir a dimensão espacial da imagem de (7,7) para (1,1) ?