(self, x)
| 390 | # return pt_checkpoint(self._forward, x) # pytorch |
| 391 | |
| 392 | def _forward(self, x): |
| 393 | b, c, *spatial = x.shape |
| 394 | x = x.reshape(b, c, -1) |
| 395 | qkv = self.qkv(self.norm(x)) |
| 396 | h = self.attention(qkv) |
| 397 | h = self.proj_out(h) |
| 398 | return (x + h).reshape(b, c, *spatial) |
| 399 | |
| 400 | |
| 401 | def count_flops_attn(model, _x, y): |