(self, input: Tensor)
| 115 | self.weight[self.padding_idx].fill_(0) |
| 116 | |
| 117 | def forward(self, input: Tensor) -> Tensor: |
| 118 | emb = F.embedding( |
| 119 | input, |
| 120 | self.weight, |
| 121 | self.padding_idx, |
| 122 | self.max_norm, |
| 123 | self.norm_type, |
| 124 | self.scale_grad_by_freq, |
| 125 | self.sparse, |
| 126 | ) |
| 127 | |
| 128 | # always apply layer norm in full precision |
| 129 | emb = emb.to(torch.get_default_dtype()) |
| 130 | |
| 131 | return self.norm(emb).to(self.weight.dtype) |
| 132 | |
| 133 | |
| 134 | class Embedding(torch.nn.Embedding): |