(self, latent_dim, time_embed_dim, dropout)
| 60 | class StylizationBlock(nn.Module): |
| 61 | |
| 62 | def __init__(self, latent_dim, time_embed_dim, dropout): |
| 63 | super().__init__() |
| 64 | self.emb_layers = nn.Sequential( |
| 65 | nn.SiLU(), |
| 66 | nn.Linear(time_embed_dim, 2 * latent_dim), |
| 67 | ) |
| 68 | self.norm = nn.LayerNorm(latent_dim) |
| 69 | self.out_layers = nn.Sequential( |
| 70 | nn.SiLU(), |
| 71 | nn.Dropout(p=dropout), |
| 72 | zero_module(nn.Linear(latent_dim, latent_dim)), |
| 73 | ) |
| 74 | |
| 75 | def forward(self, h, emb): |
| 76 | """ |
nothing calls this directly
no test coverage detected