import torch
import torch.nn.functional as F

class Model(torch.nn.Module):

    def __init__(self):
        super(Model, self).__init__()

    def forward(self, input_image: torch.Tensor, grid: torch.Tensor) -> torch.Tensor:
        return F.grid_sample(input_image, grid, mode='bilinear', padding_mode='zeros', align_corners=True)

def get_inputs():
    # 创建测试数据
    N, C, H_in, W_in = 1, 3, 256, 256
    H_out, W_out = 512, 512
    
    input_image = torch.randn(N, C, H_in, W_in)
    
    grid_y, grid_x = torch.meshgrid(
        torch.linspace(-1, 1, H_out),
        torch.linspace(-1, 1, W_out),
        indexing='ij'
    )
    grid = torch.stack((grid_x, grid_y), dim=-1) # Shape: (H_out, W_out, 2)
    grid = grid.unsqueeze(0).repeat(N, 1, 1, 1) # Shape: (N, H_out, W_out, 2)
    
    return [input_image, grid]

def get_init_inputs():
    return [] 
