Method
__init__
(self, embed_dim, n_classes=1000, key='class', ucg_rate=0.1)
Source from the content-addressed store, hash-verified
| 24 | |
| 25 | class ClassEmbedder(nn.Module): |
| 26 | def __init__(self, embed_dim, n_classes=1000, key='class', ucg_rate=0.1): |
| 27 | super().__init__() |
| 28 | self.key = key |
| 29 | self.embedding = nn.Embedding(n_classes, embed_dim) |
| 30 | self.n_classes = n_classes |
| 31 | self.ucg_rate = ucg_rate |
| 32 | |
| 33 | def forward(self, batch, key=None, disable_dropout=False): |
| 34 | if key is None: |
Tested by
no test coverage detected