File size: 385 Bytes
2e30a77 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 | import torch.nn as nn
class TokenEmbedding(nn.Module):
def __init__(
self,
vocab_size,
embedding_dim
):
super().__init__()
self.embedding = nn.Embedding(
vocab_size,
embedding_dim
)
def forward(self, input_ids):
return self.embedding(
input_ids
) |