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


BATCH_SIZE = 8
CHANNELS = 32
D_IN = 32
H_IN = 64
W_IN = 64


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

D_OUT = D_IN + PAD_F + PAD_K
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__()
        self.pad_layer = nn.CircularPad3d(padding)

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


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


def get_init_inputs():
    return [PADDING]