Update relative_transformer.py

fix typo
fastnlp · Jun 14, 2020 · 7a3859e · 7a3859e
1 parent a287270
commit 7a3859e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/modules/relative_transformer.py b/modules/relative_transformer.py
@@ -125,8 +125,8 @@ def forward(self, x, mask):
         D_ = torch.einsum('nd,ld->nl', self.r_w_bias, pos_embed)[None, :, None]  # head x 2max_len, 每个head对位置的bias
         B_ = torch.einsum('bnqd,ld->bnql', q, pos_embed)  # bsz x head  x max_len x 2max_len，每个query对每个shift的偏移
         E_ = torch.einsum('bnqd,ld->bnql', k, pos_embed)  # bsz x head x max_len x 2max_len, key对relative的bias
-        BDE = B_ + D_  # bsz x head x max_len x 2max_len, 要转换为bsz x head x max_len x max_len
-        BD = self._shift(BD) + self._transpose_shift(E_)
+        BD = B_ + D_  # bsz x head x max_len x 2max_len, 要转换为bsz x head x max_len x max_len
+        BDE = self._shift(BD) + self._transpose_shift(E_)
         attn = AC + BDE
 
         attn = attn / self.scale