(self, x: torch.tensor)
| 45 | self.dropout = nn.Dropout(dropout) |
| 46 | |
| 47 | def forward(self, x: torch.tensor): |
| 48 | batch_size = x.shape[0] |
| 49 | attentions = self.multihead_attn( |
| 50 | query=self.query.repeat(batch_size, 1, 1), |
| 51 | key=x, |
| 52 | value=x, |
| 53 | average_attn_weights=False, |
| 54 | )[0] |
| 55 | x = self.layernorm(self.dropout(attentions)) |
| 56 | return x, attentions[1] |
| 57 | |
| 58 | |
| 59 | class MLP(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected