| 403 | return x |
| 404 | |
| 405 | class TransformerBase(nn.Module): |
| 406 | def __init__(self, width, layers, heads, window_size, token_len, block_cls, drop_path_rate=0.0): |
| 407 | super().__init__() |
| 408 | self.layernorm1 = LayerNorm(width) |
| 409 | self.transformer = Transformer(width, layers, heads, window_size=window_size, block_cls=block_cls, drop_path_rate=drop_path_rate) |
| 410 | self.layernorm2 = LayerNorm(width) |
| 411 | |
| 412 | def set_grad_checkpointing(self, set_checkpointing=True): |
| 413 | self.transformer.set_grad_checkpointing(set_checkpointing) |
| 414 | |
| 415 | def forward(self, x, condition=None): |
| 416 | # x [B, V*N, D] |
| 417 | x = self.layernorm1(x) |
| 418 | x = self.transformer(x, condition) |
| 419 | x = self.layernorm2(x) |
| 420 | return x |
| 421 | |
| 422 | class TransformerEncoder(TransformerBase): |
| 423 | def __init__(self, input_res, in_channels, patch_size, width, layers, heads, window_size): |
nothing calls this directly
no outgoing calls
no test coverage detected