import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self, val=-1e9):
        super().__init__()
        self.val = val

    def forward(self, x: torch.Tensor) -> torch.Tensor:
      
        mask = torch.triu(torch.ones_like(x), diagonal=1)
        return x.masked_fill(mask.bool(), self.val)


batch_size = 16
heads = 32
seq_len = 1024 
shape = (batch_size, heads, seq_len, seq_len)

def get_inputs():
    x = torch.randn(shape, dtype=torch.float32)
    return [x]

def get_init_inputs():
    return []