import torch 
import torch.nn as nn 


class Extractor(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 4, 5)
        self.bn = nn.BatchNorm2d(4)
    
    def forward(self, x):
        return self.bn(self.conv)


class MyModel(nn.Module):
    def __init__(self, feature_extractor):
        super().__init__()
        self.fc1 = nn.Linear(16, 10)
        self.fc2 = nn.Linear(10, 10)
        self.fc3 = nn.Linear(10, 4)
        self.feature_extractor = feature_extractor 
        
    def forward(self, x):
        x = F.relu(self.feature_extractor(x))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.softmax(self.fc3(x), dim=1)
        return x


def main():
    extractor = Extractor()
    # ... train extractor 
    print(extractor.conv.weight)
    torch.save(extractor.state_dict(), "params.pth")
    
    model = MyModel(extractor)  # or MyModel(Extarctor())
    state_dict = torch.load("params.pth")
    model.feature_extractor.load_state_dict(state_dict)
    print(model.feature_extractor.conv.weight)

main()