(self, x)
| 319 | #return pt_checkpoint(self._forward, x) # pytorch |
| 320 | |
| 321 | def _forward(self, x): |
| 322 | b, c, *spatial = x.shape |
| 323 | x = x.reshape(b, c, -1) |
| 324 | qkv = self.qkv(self.norm(x)) |
| 325 | h = self.attention(qkv) |
| 326 | h = self.proj_out(h) |
| 327 | return (x + h).reshape(b, c, *spatial) |
| 328 | |
| 329 | |
| 330 | def count_flops_attn(model, _x, y): |
nothing calls this directly
no outgoing calls
no test coverage detected