(self, query_size, key_size, proj_size)
| 65 | @registry.register('pointer', 'bahdanau') |
| 66 | class BahdanauPointer(torch.nn.Module): |
| 67 | def __init__(self, query_size, key_size, proj_size): |
| 68 | super().__init__() |
| 69 | self.compute_scores = torch.nn.Sequential( |
| 70 | torch.nn.Linear(query_size + key_size, proj_size), |
| 71 | torch.nn.Tanh(), |
| 72 | torch.nn.Linear(proj_size, 1)) |
| 73 | |
| 74 | def forward(self, query: torch.Tensor, keys: torch.Tensor, attn_mask=None): |
| 75 | # query shape: batch x query_size |