c-my/transformers
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

improve comment

Browse Source
This commit is contained in:
Patrick von Platen
2020-04-22 07:55:19 +02:00
parent d911f44507
commit 598caff988
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/transformers/modeling_reformer.py
+2 -2
View File
@@ -218,12 +218,12 @@ class LSHSelfAttention(nn.Module):
# Causal mask
if self.is_decoder:
# TODO: This line can be improved in terms of memory
# TODO (PVP): This line can be improved in terms of memory. Mask should be inserted and does not have to be created each time layer is called
mask = torch.lt(query_info.unsqueeze(-1), key_value_info.unsqueeze(-2)).long().to(query_info.device)
query_key_dots = query_key_dots - mask * 1e9
# Self mask
# TODO: This line can be improved in terms of memory
# TODO (PVP): This line can be improved in terms of memory. Mask should be inserted and does not have to be created each time layer is called
mask = torch.eq(query_info.unsqueeze(-1), key_value_info.unsqueeze(-2)).long().to(query_info.device)
query_key_dots = query_key_dots - mask * 1e5