| 23 | |
| 24 | |
| 25 | class AbsolutePositionalEmbedding(nn.Module): |
| 26 | def __init__(self, dim, max_seq_len): |
| 27 | super().__init__() |
| 28 | self.emb = nn.Embedding(max_seq_len, dim) |
| 29 | self.init_() |
| 30 | |
| 31 | def init_(self): |
| 32 | nn.init.normal_(self.emb.weight, std=0.02) |
| 33 | |
| 34 | def forward(self, x): |
| 35 | n = torch.arange(x.shape[1], device=x.device) |
| 36 | return self.emb(n)[None, :, :] |
| 37 | |
| 38 | |
| 39 | class FixedPositionalEmbedding(nn.Module): |