Method __init__

tensorrt_llm/layers/embedding.py:194–207

(self,
                 num_classes: int,
                 hidden_size: int,
                 dropout_prob: float = 0.0,
                 mapping=Mapping(),
                 dtype=None)

class LabelEmbedding(Module):

    def __init__(self,
                 num_classes: int,
                 hidden_size: int,
                 dropout_prob: float = 0.0,
                 mapping=Mapping(),
                 dtype=None):
        super().__init__()
        use_cfg_embedding = dropout_prob > 0
        self.embedding_table = Embedding(num_classes + use_cfg_embedding,
                                         hidden_size,
                                         tp_size=mapping.tp_size,
                                         tp_group=mapping.tp_group,
                                         dtype=dtype)
        self.num_classes = num_classes

    def token_drop(self, labels: Tensor, force_drop_ids: Tensor):
        labels = where(force_drop_ids == 1, self.num_classes, labels)