(self, x, cond)
| 54 | ) |
| 55 | |
| 56 | def forward(self, x, cond): |
| 57 | # x: [N, L, D] |
| 58 | # cond: [N, L_cond, D_cond] |
| 59 | x = x + self.cross_attn(self.norm1(x), cond, cond)[0] |
| 60 | before_sa = self.norm2(x) |
| 61 | x = x + self.self_attn(before_sa, before_sa, before_sa)[0] |
| 62 | x = x + self.mlp(self.norm3(x)) |
| 63 | return x |
| 64 | |
| 65 | |
| 66 | class TriplaneTransformer(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected