| 2 | |
| 3 | |
| 4 | class TriangularCausalMask(): |
| 5 | def __init__(self, B, L, device="cpu"): |
| 6 | mask_shape = [B, 1, L, L] |
| 7 | with torch.no_grad(): |
| 8 | self._mask = torch.triu(torch.ones(mask_shape, dtype=torch.bool), diagonal=1).to(device) |
| 9 | |
| 10 | @property |
| 11 | def mask(self): |
| 12 | return self._mask |
| 13 | |
| 14 | |
| 15 | class ProbMask(): |