(self, encoder_dim: int, time_embed_dim: int, num_heads: int = 64)
| 1885 | |
| 1886 | class TextTimeEmbedding(nn.Module): |
| 1887 | def __init__(self, encoder_dim: int, time_embed_dim: int, num_heads: int = 64): |
| 1888 | super().__init__() |
| 1889 | self.norm1 = nn.LayerNorm(encoder_dim) |
| 1890 | self.pool = AttentionPooling(num_heads, encoder_dim) |
| 1891 | self.proj = nn.Linear(encoder_dim, time_embed_dim) |
| 1892 | self.norm2 = nn.LayerNorm(time_embed_dim) |
| 1893 | |
| 1894 | def forward(self, hidden_states): |
| 1895 | hidden_states = self.norm1(hidden_states) |
nothing calls this directly
no test coverage detected