Drops labels to enable classifier-free guidance.
(self, labels, force_drop_ids=None)
| 1522 | self.dropout_prob = dropout_prob |
| 1523 | |
| 1524 | def token_drop(self, labels, force_drop_ids=None): |
| 1525 | """ |
| 1526 | Drops labels to enable classifier-free guidance. |
| 1527 | """ |
| 1528 | if force_drop_ids is None: |
| 1529 | drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob |
| 1530 | else: |
| 1531 | drop_ids = torch.tensor(force_drop_ids == 1) |
| 1532 | labels = torch.where(drop_ids, self.num_classes, labels) |
| 1533 | return labels |
| 1534 | |
| 1535 | def forward(self, labels: torch.LongTensor, force_drop_ids=None): |
| 1536 | use_dropout = self.dropout_prob > 0 |