import torch
import torch.nn as nn
import torch.nn.functional as F

BATCH = 64
CHANNELS = 3
H_IN, W_IN = 512, 512
H_OUT, W_OUT = 128, 128
BOX = [100, 100, 400, 400] # y1, x1, y2, x2

class Model(nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
 
        return F.interpolate(x[..., BOX[0]:BOX[2], BOX[1]:BOX[3]], size=(H_OUT, W_OUT), mode='bilinear', align_corners=False)

def get_inputs():
    x = torch.rand(BATCH, CHANNELS, H_IN, W_IN, device='cuda', dtype=torch.float32)
    return [x]

def get_init_inputs():
    return []