| 508 | return x + skip |
| 509 | |
| 510 | class GlobalTransformerLayer(nn.Module): |
| 511 | def __init__(self, d_model, d_ff, d_head, cond_features, dropout=0.0): |
| 512 | super().__init__() |
| 513 | self.self_attn = SelfAttentionBlock(d_model, d_head, cond_features, dropout=dropout) |
| 514 | self.ff = FeedForwardBlock(d_model, d_ff, cond_features, dropout=dropout) |
| 515 | |
| 516 | def forward(self, x, pos, cond): |
| 517 | x = checkpoint(self.self_attn, x, pos, cond) |
| 518 | x = checkpoint(self.ff, x, cond) |
| 519 | return x |
| 520 | |
| 521 | class NeighborhoodTransformerLayer(nn.Module): |
| 522 | def __init__(self, d_model, d_ff, d_head, cond_features, kernel_size, dropout=0.0): |