I defining two model which are encoder and classfication respectively. I want use encoder’s ouput as the input of the classfication then train the two model. After training I want to save the two model separately.
here’s my two model
# encoder model 
class source(nn.Module):
    def __init__(self, initial_channel):
        super(source, self).__init__()
        self.initial_channel = initial_channel


        def conv_bn(inp, oup, stride):
            return nn.Sequential(
                nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
                nn.BatchNorm2d(oup),
                nn.ReLU(inplace=True)
            )
        self.model = self.model = nn.Sequential(conv_bn(initial_channel, 64, 2),
                                   conv_bn(64, 128, 2),
                                   conv_bn(128, 128, 1),
                                   conv_bn(128, 256, 2),
                                   conv_bn(256, 256, 1),
                                   nn.AdaptiveAvgPool2d([1,1]))
    def forward(self, x):
        x = self.model(x)
        x = x.view(-1,256)
        return x

# classfication model
class classfication(nn.Module):
    def __init__(self, class_num):
        super(classfication, self).__init__()
        self.class_num = class_num

        self.fc1 = nn.Sequential(nn.Linear(256, 64),
                                 nn.ReLU(inplace=True),
                                 nn.Dropout(0.5))
        self.fc2 = nn.Sequential(nn.Linear(64, class_num), )

    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        return x

And following my main model to describe what I want to do in detail.
encoder = source(initial_channel)
cls = classfication(classnum)

model = torch.nn.Sequential(encoder, cls)
train(model)
save(encoder)
save(cls)