simpler subsequent mask generator (#1198)

pytorch · Nov 27, 2023 · c0b889d · c0b889d
1 parent c4dc481
commit c0b889d
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/word_language_model/model.py b/word_language_model/model.py
@@ -120,9 +120,7 @@ def __init__(self, ntoken, ninp, nhead, nhid, nlayers, dropout=0.5):
         self.init_weights()
 
     def _generate_square_subsequent_mask(self, sz):
-        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
-        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
-        return mask
+        return torch.log(torch.tril(torch.ones(sz,sz)))
 
     def init_weights(self):
         initrange = 0.1