h: B, T, D emb: B, D
(self, h, emb)
| 73 | ) |
| 74 | |
| 75 | def forward(self, h, emb): |
| 76 | """ |
| 77 | h: B, T, D |
| 78 | emb: B, D |
| 79 | """ |
| 80 | # B, 1, 2D |
| 81 | emb_out = self.emb_layers(emb).unsqueeze(1) |
| 82 | # scale: B, 1, D / shift: B, 1, D |
| 83 | scale, shift = torch.chunk(emb_out, 2, dim=2) |
| 84 | h = self.norm(h) * (1 + scale) + shift |
| 85 | h = self.out_layers(h) |
| 86 | return h |
| 87 | |
| 88 | |
| 89 | class LinearTemporalSelfAttention(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected