import torch
import torch.nn as nn

torch.backends.cuda.matmul.allow_tf32 = False

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()

    def forward(self, indices: torch.Tensor) -> torch.Tensor:
  
        return (indices != -1).float()

N = 1024 * 1024 * 32 

def get_inputs():

    indices = torch.randint(0, 1000, (N,), dtype=torch.int32, device='cuda')

    mask = torch.rand(N, device='cuda') < 0.1
    indices[mask] = -1
    
    return [indices]

def get_init_inputs():
    return []