| 487 | |
| 488 | class TransformerDecoder(TransformerBase): |
| 489 | def __init__(self, token_len, width, layers, heads, window_size, encoder_dim=None): |
| 490 | self.width = width |
| 491 | super().__init__(width, layers, heads, window_size, token_len, ResAttBlock) |
| 492 | self.positional_embedding = nn.Parameter(torch.zeros(1, token_len, width)) |
| 493 | nn.init.trunc_normal_(self.positional_embedding, std=0.02) |
| 494 | |
| 495 | if encoder_dim is not None and encoder_dim != width: |
| 496 | self.encoder_proj = nn.Linear(encoder_dim, width) |
| 497 | self.out_proj = nn.Linear(width, encoder_dim) |
| 498 | else: |
| 499 | self.encoder_proj = nn.Identity() |
| 500 | self.out_proj = nn.Identity() |
| 501 | |
| 502 | |
| 503 | def forward(self, latent, condition=None, reverse=False): |