如何在Huggingface BERT模型之上添加LSTM层

import torch
import torch.nn as nn
from transformers import BertModel

class BertLstmModel(nn.Module):
    def __init__(self):
        super(BertLstmModel, self).__init__()
        
        self.bert = BertModel.from_pretrained('bert-base-uncased')
        self.lstm = nn.LSTM(input_size=768, hidden_size=256, num_layers=2, batch_first=True)
        self.fc = nn.Linear(256, num_classes)  # num_classes是分类任务的类别数量

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        sequence_output = outputs.last_hidden_state
        
        lstm_output, _ = self.lstm(sequence_output)
        lstm_output = lstm_output[:, -1, :]  # 取最后一个时刻的输出
        
        logits = self.fc(lstm_output)
        return logits

model = BertLstmModel()
model.load_state_dict(torch.load('bert_lstm_model.pth'))
model.eval()

input_ids = torch.tensor([input_ids])  # input_ids是输入序列的索引
attention_mask = torch.tensor([attention_mask])  # attention_mask用于掩盖填充部分
logits = model(input_ids, attention_mask)
predictions = torch.argmax(logits, dim=1)