- mask = torch.tril(q.new_ones(a.size(2), a.size(3)))[None, None, :, :] == 0
- a = a.masked_fill(mask, float('-inf'))
+ forbidden_attention = torch.arange(a.size(2), device = q.device)[None, None, :, None] \
+ < torch.arange(a.size(3), device = q.device)[None, None, None, :]
+ a = a.masked_fill(forbidden_attention, float('-inf'))
+