Hi,
I am working on a problem that requires pre-training a first model at the beginning and then using this pre-trained model and fine-tuning it along with a second model. When training the first model, it requires a classification layer in order to compute a loss for it. However, I do not need my classification layer when using the pretrained model along with my second model. I only need the output  (which in my case is the hidden state of an LSTM). At the second stage, when loading the pre-trained model, if i remove the classification layer, then does PyTorch automatically ignore the weights for that classification layer and keep the rest when fine-tuning? Are the weights for each sub-module separately saved and then used as needed? This is a snippet of the code:
class LSTMD(nn.Module):
    
    def __init__(self, input_size, hidden_size, num_classes):
        super(LSTMD, self).__init__()
        self.hidde_size = hidden_size
        self.embedding = nn.Embedding(num_classes, hidden_size)
        self.lstm = nn.LSTMCell(input_size, hidden_size)
        self.classification = nn.Linear(hidden_size, num_classes)
        
    def forward(self, x):
        h,c = [torch.zeros(1, self.hidden_size), torch.zeros(1, self.hidden_size)]
        x = self.embedding(x)
        h,c = self.lstm(x, (h,c))
        out = self.classification(h)
        return out

model = LSTMD(512,512,1000)
state = {'model': model, 'model_optimizer': model_optimizer}

torch.save(state, 'saved.pth')

After that, I only want to load my pre-trained model but only use h from the LSTM output and discard the classification layer when fine-tuning the model.
class New(nn.Module):
    
    def __init__(self, checkpoint):
        super(New, self).__init__()
        checkpoint = torch.load(checkpoint)
        old_model = checkpoint['model']
        # I removed out = self.classification(h) from the old model
        modules = list(old_model.children())[:-1]
        self.new_model = nn.Sequential(*modules)

    def forward(self, x):

        out = self.new_model(x)  
        return out


But then I get the error: forward() takes _ positional arguments but _ were given
_ depends on the arguments in my case.
However, when I remove the
self.classification = nn.Linear(hidden_size, num_classes)
and
out = self.classification(h)
from the first model, and load the checkpoint, I get no error. But i’m afraid this is not the correct way as later the model will have problem fine-tuning? The problem mainly happens when I call the forward function on the new model created by nn.Sequential.  Even if i keep the layers by doing:
 modules = list(old_model.children())
  self.new_model = nn.Sequential(*modules)

And then run the forward function, it dosen’t work. Something is happening is the sequential operation .Any clue what may be the problem
Thanks!