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


BATCH_SIZE = 128
CHANNELS = 64
H_IN = 64
W_IN = 64

PADDING = (2, 3, 4, 5) 
PAD_L, PAD_R, PAD_T, PAD_B = PADDING

PAD_VAL = 1.5

H_OUT = H_IN + PAD_T + PAD_B
W_OUT = W_IN + PAD_L + PAD_R

class Model(nn.Module):
   

    def __init__(self, padding, value):
        super().__init__()
        self.pad_layer = nn.ConstantPad2d(padding, value)

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


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


def get_init_inputs():
    return [PADDING, PAD_VAL]