(self, num_classes, hidden_size, dropout_prob)
| 89 | dropout for classifier-free guidance. |
| 90 | """ |
| 91 | def __init__(self, num_classes, hidden_size, dropout_prob): |
| 92 | super().__init__() |
| 93 | use_cfg_embedding = int(dropout_prob > 0) |
| 94 | self.embedding_table = ParallelEmbedding( |
| 95 | num_classes + use_cfg_embedding, hidden_size, |
| 96 | init_method=functools.partial(nn.init.normal_, std=0.02), |
| 97 | ) |
| 98 | self.num_classes = num_classes |
| 99 | self.dropout_prob = dropout_prob |
| 100 | |
| 101 | def token_drop(self, labels, force_drop_ids=None): |
| 102 | """ |