Method
__init__
(self, in_channels: int, cond_embed_dim: int, act_fn: str = "silu")
Source from the content-addressed store, hash-verified
| 43 | |
| 44 | class CondEmbedding(nn.Module): |
| 45 | def __init__(self, in_channels: int, cond_embed_dim: int, act_fn: str = "silu"): |
| 46 | super().__init__() |
| 47 | |
| 48 | self.linear = nn.Linear(in_channels, cond_embed_dim) |
| 49 | self.act = None |
| 50 | if act_fn == "silu": |
| 51 | self.act = nn.SiLU() |
| 52 | elif act_fn == "mish": |
| 53 | self.act = nn.Mish() |
| 54 | |
| 55 | def forward(self, sample): |
| 56 | sample = self.linear(sample) |
Tested by
no test coverage detected