model/embedding.py

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
        )