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, 4, 3, 5) 
PAD_L, PAD_R, PAD_T, PAD_B = PADDING

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

class Model(nn.Module):
  
    def __init__(self, padding):
        super().__init__()
        # padding: (pad_left, pad_right, pad_top, pad_bottom)
        self.pad_layer = nn.CircularPad2d(padding)

    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]