Hi. I’m trying to create a model by using sub_layers of two different models. the first model is defined in this way:
class Model5Layer_128(nn.Module):
    def __init__(self):
        super(Model5Layer_128, self).__init__()
        self.conv1 = nn.Conv2d(in_channels= 3,out_channels= 6,kernel_size= 5)
        self.pool = nn.MaxPool2d(kernel_size=2,stride=2)
        self.conv2 = nn.Conv2d(in_channels= 6, out_channels= 16, kernel_size= 5)
        self.fc1 = nn.Linear(in_features=16 * 5 * 5, out_features= 128)
        self.fc2 = nn.Linear(in_features= 128,out_features= 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.softmax(x,dim=-1)
        return x

Then I use a fraction of an instance of this model + a regressor and train it. For the second phase, I remove the regressor and try to reimplement the model( a model with the same architecture of the model above) by using this method:
# remove the regressor from our model, return the trained partial model
def get_cnn_regressor_removed_partial_student():
    basic_model = Model5Layer_128()
    cnn_regressor_added_student = nn.Sequential(*list(basic_model.children())[:3])
    return cnn_regressor_added_student

combined_student = nn.Sequential(regressor_removed_student,*list(model_128.children())[3:])
print(combined_student)

It seems correct, but when I try to train the whole model or use torch summary, I get this error
RuntimeError: size mismatch, m1: [320 x 10], m2: [400 x 128] at /pytorch/aten/src/TH/generic/THTensorMath.cpp:41

Any help would be appreciated.