tensorflowとpytorch間でパラメータ数が合わない

tensorflowで構築したCNNモデルをpytorchでかき直しているのですがどうもうまくいきません。
入力は(40, 40, 4)の画像です。

以下、tensorflowのコードです。

from keras.models import Sequential, Model
from keras.layers import Conv2D, MaxPooling2D, Input
from keras.layers.core import Dense, Dropout, Activation, Flatten

model = Sequential()
model.add(Conv2D(64, 3, 3, activation='relu', input_shape=input_shape))
model.add(Conv2D(64, 3, 3, activation='relu'))
model.add(MaxPooling2D((2,2), strides=(2,2)))

model.add(Conv2D(128, 3, 3, activation='relu'))
model.add(Conv2D(128, 3, 3, activation='relu'))
model.add(MaxPooling2D((2,2), strides=(2,2)))

model.add(Flatten())
model.add(Dense(256, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(128, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(1, activation='sigmoid'))

Flattenした後の次元は6272です。

一方でpytorchでは

import torch.nn as nn

class Model(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(4, 64, kernel_size=3, stride=1, padding=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.layer2 = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=2),
            nn.ReLU(),
            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.fc = nn.Sequential(
            nn.Linear(21632, 256),
            nn.ReLU(),
            nn.Dropout(p=0.5),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Dropout(p=0.5),
            nn.Linear(128, 1))
        
    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc(out)
      
        return out

こちらは畳み込み層の後flattenした時に21632もあります。

padding等のズレがあるにしても三倍もの差が生まれるのはどうしてかなと思って質問させていただきました。

ちなみにtensorflowの方は無事lossが下がるのですが、pytorchの方は下がらず予測値がある一定の値の周辺に固まっています。もちろん、ネットワークの構造以外の部分などの他の要因も考えられるのですが、まずはこのパラメータ数の違いをなんとかしたいと思い質問させていただきました。

ご教授よろしくお願いします。