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, input: torch.Tensor, index: torch.Tensor) -> torch.Tensor:

        return torch.gather(input, dim=-1, index=index.long())


rows = 1024 * 16
cols = 256

def get_inputs():

    input = torch.randn(rows, cols, dtype=torch.float32, device='cuda').contiguous()

    index = torch.randint(0, cols, (rows, cols), dtype=torch.int32, device='cuda').contiguous()
    
    return [input, index]

def get_init_inputs():
    return []