(self, x, **kwargs)
| 166 | return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b) |
| 167 | |
| 168 | def forward(self, x, **kwargs): |
| 169 | h_ = x |
| 170 | h_ = self.attention(h_) |
| 171 | h_ = self.proj_out(h_) |
| 172 | return x + h_ |
| 173 | |
| 174 | |
| 175 | class MemoryEfficientAttnBlock(nn.Module): |