import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs 2D convolution operation.
    """
    def __init__(self, weight, bias=None):
        super(Model, self).__init__()
        self.weight = nn.Parameter(weight)
        self.bias = nn.Parameter(bias) if bias is not None else None
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs 2D convolution.

        Args:
            x (torch.Tensor): Input tensor of shape [batch_size, in_channels, height, width]

        Returns:
            torch.Tensor: Output tensor of shape [batch_size, out_channels, out_height, out_width]
        """
        return torch.nn.functional.conv2d(x, self.weight, self.bias, stride=1, padding=0)

# Hyperparameters
batch_size = 4
in_channels = 3
out_channels = 64
height = 32
width = 32
kernel_size = 3

def get_inputs():
    x = torch.randn(batch_size, in_channels, height, width)
    return [x]

def get_init_inputs():
    weight = torch.randn(out_channels, in_channels, kernel_size, kernel_size)
    bias = torch.randn(out_channels)
    return [weight, bias]