(self, num_buckets, num_heads, bidirectional, max_dist=128)
| 147 | class T5RelativeEmbedding(nn.Module): |
| 148 | |
| 149 | def __init__(self, num_buckets, num_heads, bidirectional, max_dist=128): |
| 150 | super(T5RelativeEmbedding, self).__init__() |
| 151 | self.num_buckets = num_buckets |
| 152 | self.num_heads = num_heads |
| 153 | self.bidirectional = bidirectional |
| 154 | self.max_dist = max_dist |
| 155 | |
| 156 | # layers |
| 157 | self.embedding = nn.Embedding(num_buckets, num_heads) |
| 158 | |
| 159 | def forward(self, lq, lk): |
| 160 | device = self.embedding.weight.device |