| 411 | return x |
| 412 | |
| 413 | class LayerNormProxy(nn.Module): |
| 414 | |
| 415 | def __init__(self, dim): |
| 416 | |
| 417 | super().__init__() |
| 418 | self.norm = nn.LayerNorm(dim) |
| 419 | |
| 420 | def forward(self, x): |
| 421 | |
| 422 | x = einops.rearrange(x, 'b c h w -> b h w c') |
| 423 | x = self.norm(x) |
| 424 | return einops.rearrange(x, 'b h w c -> b c h w') |
| 425 | |
| 426 | |
| 427 | class TransformerMLPWithConv(nn.Module): |