I have a 2-layer bidirectional LSTM, and I need to remove the second layer of the pre-trained LSTM. My model is very simple and looks like this:

num_layers = 2
class BiRNN(nn.Module):
def init(self, input_size, hidden_size, num_layers, num_classes):
super(BiRNN, self).init()
self.hidden_size = hidden_size
self.num_layers = num_layers
self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, bidirectional=True)
self.fc = nn.Linear(hidden_size2, num_classes)
def forward(self, x):
# Set initial states
h0 = torch.zeros(self.num_layers2, x.size(0), self.hidden_size).to(device)
c0 = torch.zeros(self.num_layers*2, x.size(0), self.hidden_size).to(device)
    # Forward propagate LSTM
    lstm_out, _ = self.lstm(x, (h0, c0))  # lstm_out: tensor of shape (batch_size, seq_length, hidden_size*2)

    # Decode the hidden state of the last time step
    out = self.fc(lstm_out[:, -1, :])
    return out

initialize the NN
model = BiRNN(input_size, hidden_size, num_layers, num_classes).to(device)

I train the network, and now I should remove the second layer of the LSTM for some analysis. I checked this post https://discuss.pytorch.org/t/can-i-remove-a-layer-from-a-pre-trained-model-while-loading-the-model-weights/57899 5. But in that post, the goal is to remove the fully connected layer. I tried to use model.childern, but both layers of LSTM are counted as one child. The same is true with model.modules. I also read about nn.Identity, but not sure how to replace the second layer of LSTM by nn.Identity. Is there any way to remove the second layer of the model in PyTorch.
Sorry if the question is trivial, I am new to PyTorch.
Thanks!
> Blockquote