import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Simple model that performs InstanceNorm operation.
    """

    def __init__(self, num_features=64, eps=1e-5, affine=True, track_running_stats=False):
        super(Model, self).__init__()
        self.num_features = num_features
        self.eps = eps
        self.affine = affine
        self.track_running_stats = track_running_stats

        # 创建InstanceNorm层
        self.instance_norm = nn.InstanceNorm2d(
            num_features=num_features,
            eps=eps,
            affine=affine,
            track_running_stats=track_running_stats
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies InstanceNorm to the input tensor.

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

        Returns:
            torch.Tensor: Output tensor after instance normalization, same shape as input.
        """
        return self.instance_norm(x)


# 参数配置
batch_size = 32
num_features = 64
height = 128
width = 128


def get_inputs():
    """
    生成InstanceNorm的输入张量。

    Returns:
        list: 包含一个形状为 [batch_size, num_features, height, width] 的张量
    """
    x = torch.randn(batch_size, num_features, height, width)
    return [x]


def get_init_inputs():
    """
    获取模型初始化所需的输入（空列表，因为不需要特殊初始化）。

    Returns:
        list: 空列表
    """
    return []  # No special initialization inputs needed
