| 494 | |
| 495 | |
| 496 | class TextTimeEmbedding(nn.Module): |
| 497 | def __init__(self, encoder_dim: int, time_embed_dim: int, num_heads: int = 64): |
| 498 | super().__init__() |
| 499 | self.norm1 = nn.LayerNorm(encoder_dim) |
| 500 | self.pool = AttentionPooling(num_heads, encoder_dim) |
| 501 | self.proj = nn.Linear(encoder_dim, time_embed_dim) |
| 502 | self.norm2 = nn.LayerNorm(time_embed_dim) |
| 503 | |
| 504 | def forward(self, hidden_states): |
| 505 | hidden_states = self.norm1(hidden_states) |
| 506 | hidden_states = self.pool(hidden_states) |
| 507 | hidden_states = self.proj(hidden_states) |
| 508 | hidden_states = self.norm2(hidden_states) |
| 509 | return hidden_states |
| 510 | |
| 511 | |
| 512 | class TextImageTimeEmbedding(nn.Module): |
no outgoing calls
no test coverage detected