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, params: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
       
        b = indices[:, 0]
        h = indices[:, 1]
        w = indices[:, 2]
        
        return params[b, h, w]

B, H, W, C = 4, 256, 256, 128
N = 1024 * 1024 

def get_inputs():
    params = torch.randn(B, H, W, C, dtype=torch.float32).cuda()

    idx_b = torch.randint(0, B, (N,)).cuda()
    idx_h = torch.randint(0, H, (N,)).cuda()
    idx_w = torch.randint(0, W, (N,)).cuda()
    
    indices = torch.stack([idx_b, idx_h, idx_w], dim=1).long() # [N, 3]
    
    return [params, indices]

def get_init_inputs():
    return []