(self, x, index, condition=None)
| 289 | return x |
| 290 | |
| 291 | def forward(self, x, index, condition=None): |
| 292 | # no condition in encoder, its a dummy argument |
| 293 | y = self.layernorm1(x) |
| 294 | y = self.attention(y, index) |
| 295 | x = x.type(torch.float32) + y # residual in fp32 |
| 296 | y = self.layernorm2(x) |
| 297 | y = self.mlp(y) |
| 298 | x = x.type(torch.float32) + y # residual in fp32 |
| 299 | return x |
| 300 | |
| 301 | class ConditionalResAttBlock(nn.Module): |
| 302 | def __init__(self, d_model, n_head, window_size=None, drop_path_rate=0.0): |