to match AttnBlock usage
| 135 | |
| 136 | |
| 137 | class LinAttnBlock(LinearAttention): |
| 138 | """to match AttnBlock usage""" |
| 139 | |
| 140 | def __init__(self, in_channels): |
| 141 | super().__init__(dim=in_channels, heads=1, dim_head=in_channels) |
| 142 | |
| 143 | |
| 144 | class AttnBlock(nn.Module): |