def forward()

in opacus_lab/models/GPT2/model/masking.py [0:0]

• 8 lines of code
• 1 McCabe index (conditional complexity)

    def forward(self, x: torch.Tensor, offset: int = 0) -> torch.Tensor:
        seq_len = x.size(-1)

        # Create shifted upper triangular matrix.
        future = torch.ones(
            (seq_len, seq_len + offset), dtype=torch.bool, device=x.device
        )
        future = future.triu(offset + 1)

        mask = future.view((1,) * (x.ndim - 1) + future.size())
        return mask.expand(x.shape + mask.shape[-1:])