import torch
import torch.nn as nn

BATCH = 16
CHANNELS = 64
LENGTH = 4096
KERNEL_SIZE = 3
DILATION = 2
PADDING = 2 
STRIDE = 1

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        torch.backends.cudnn.enabled = False
        
        self.conv = nn.Conv1d(
            in_channels=CHANNELS, 
            out_channels=CHANNELS, 
            kernel_size=KERNEL_SIZE, 
            stride=STRIDE,
            padding=PADDING, 
            dilation=DILATION,
            groups=CHANNELS, 
            bias=False
        )
        nn.init.constant_(self.conv.weight, 1.0)

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

def get_inputs():
    x = torch.randint(low=-2, high=3, size=(BATCH, CHANNELS, LENGTH), device='cuda').float()
    return [x]

def get_init_inputs():
    return []