Hello,
I have a model training setup as follows:
input data > feature extraction cnn > main prediction model.
My question is; after training the model, I save the state_dict; does this save the parameters of the feature extractor neural network in addition to the the parameters of MyModel?  Main model looks something like this:
class MyModel(nn.Module):
    def __init__(self, feature_extractor):
        super(MyModel, self).__init__()
        self.fc1 = nn.Linear(16, 10)
        self.fc2 = nn.Linear(10, 10)
        self.fc3 = nn.Linear(10, 4)
        self.feature_extractor = feature_extractor  #this is a CNN network that transforms the original data
        
    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

model = MyModel()
### train code here ###
th.save(model.state_dict())

If the answer is no, how can I ensure that parameters of both models are saved after training?
And also, what is a neat way to save/load parameters of both models?
Thank you.