Skip to content

Commit c0b889d

Browse files
authored
simpler subsequent mask generator (#1198)
1 parent c4dc481 commit c0b889d

File tree

1 file changed

+1
-3
lines changed

1 file changed

+1
-3
lines changed

word_language_model/model.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -120,9 +120,7 @@ def __init__(self, ntoken, ninp, nhead, nhid, nlayers, dropout=0.5):
120120
self.init_weights()
121121

122122
def _generate_square_subsequent_mask(self, sz):
123-
mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
124-
mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
125-
return mask
123+
return torch.log(torch.tril(torch.ones(sz,sz)))
126124

127125
def init_weights(self):
128126
initrange = 0.1

0 commit comments

Comments
 (0)