import torch
import torch.nn as nn


# Shape: (N, C, W)
BATCH_SIZE = 256
CHANNELS = 128
W_IN = 4096  

PADDING = (3, 1) 
PAD_VAL = 3.5


PAD_L, PAD_R = PADDING
W_OUT = W_IN + PAD_L + PAD_R

class Model(nn.Module):


    def __init__(self, padding, value):
        super().__init__()

        self.pad_layer = nn.ConstantPad1d(padding, value)

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


def get_inputs():

    x = torch.randn(BATCH_SIZE, CHANNELS, W_IN, dtype=torch.float32)
    return [x]


def get_init_inputs():
    return [PADDING, PAD_VAL]