帮我用bert和pytorch等价实现nn.Embedding()

可以使用以下代码实现：

import torch
from transformers import BertModel

class BertEmbedding(torch.nn.Module):
def __init__(self, bert_model):
super(BertEmbedding, self).__init__()
self.bert = bert_model

def forward(self, input_ids):
outputs = self.bert(input_ids)
pooled_output = outputs[1]
return pooled_output

# Load pre-trained BERT model
bert_model = BertModel.from_pretrained('bert-base-uncased')

# Create BertEmbedding module
bert_embedding = BertEmbedding(bert_model)

# Test with example input
input_ids = torch.tensor([[1, 2, 3, 4, 5]])
embedding = bert_embedding(input_ids)
print(embedding.shape) # Output: torch.Size([1, 768])