import torch
import torch.nn as nn

BATCH = 16
CHANNELS = 64
HEIGHT = 64
WIDTH = 64
KERNEL_SIZE = 3
STRIDE = 1
PADDING = 1

class Model(nn.Module):
    def __init__(self):
        super().__init__()
 
        self.fold = nn.Fold(
            output_size=(HEIGHT, WIDTH),
            kernel_size=KERNEL_SIZE,
            dilation=1,
            padding=PADDING,
            stride=STRIDE
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.fold(x)

def get_inputs():

    L = HEIGHT * WIDTH
    C_in = CHANNELS * KERNEL_SIZE * KERNEL_SIZE

    x = torch.randint(low=-2, high=3, size=(BATCH, C_in, L), device='cuda').float()
    return [x]

def get_init_inputs():
    return []