| 71 | |
| 72 | |
| 73 | class LabelEmbedder(Module): |
| 74 | |
| 75 | def __init__(self, num_classes, hidden_size, dropout_prob, dtype=None): |
| 76 | super().__init__() |
| 77 | use_cfg_embedding = dropout_prob > 0 |
| 78 | self.embedding_table = Embedding(num_classes + use_cfg_embedding, |
| 79 | hidden_size, |
| 80 | dtype=dtype) |
| 81 | self.num_classes = num_classes |
| 82 | self.dropout_prob = dropout_prob |
| 83 | |
| 84 | def forward(self, labels, force_drop_ids=None): |
| 85 | assert force_drop_ids is None |
| 86 | embeddings = self.embedding_table(labels) |
| 87 | return embeddings |
| 88 | |
| 89 | |
| 90 | class PatchEmbed(Module): |