(self, d_model=512, n_head=8)
| 451 | |
| 452 | class ScaledDotProductAttention(layer.Layer): |
| 453 | def __init__(self, d_model=512, n_head=8): |
| 454 | super(ScaledDotProductAttention, self).__init__() |
| 455 | self.d_k = d_model // n_head |
| 456 | assert ( |
| 457 | self.d_k * n_head == d_model |
| 458 | ), "embed_dim must be divisible by num_heads" |
| 459 | |
| 460 | def forward(self, query, key, value, attn_mask): |
| 461 | """ |