h: B, T, D emb: B, D
(self, h, emb)
| 27 | ) |
| 28 | |
| 29 | def forward(self, h, emb): |
| 30 | """ |
| 31 | h: B, T, D |
| 32 | emb: B, D |
| 33 | """ |
| 34 | # B, 1, 2D |
| 35 | emb_out = self.emb_layers(emb).unsqueeze(1) |
| 36 | # scale: B, 1, D / shift: B, 1, D |
| 37 | scale, shift = torch.chunk(emb_out, 2, dim=2) |
| 38 | h = self.norm(h) * (1 + scale) + shift |
| 39 | h = self.out_layers(h) |
| 40 | return h |
nothing calls this directly
no outgoing calls
no test coverage detected