import torch
import torch.nn as nn


# --- 用于基准测试的配置 ---
batch_size = 512
dim = 4096
DELTA = 1.0
REDUCTION = 'mean'

class Model(nn.Module):
    """
    使用 PyTorch 内置的 torch.nn.HuberLoss 作为基准模型。
    """
    def __init__(self, delta=1.0, reduction='mean'):
        super(Model, self).__init__()
        self.loss_fn = nn.HuberLoss(delta=delta, reduction=reduction)
    
    def forward(self, input_tensor: torch.Tensor, target_tensor: torch.Tensor) -> torch.Tensor:
        return self.loss_fn(input_tensor, target_tensor)

def get_inputs():
    """
    生成用于测试的输入张量。
    """
    input_tensor = torch.randn(batch_size, dim, dtype=torch.float32)
    target_tensor = input_tensor + torch.randn(batch_size, dim, dtype=torch.float32) * 0.5
    return [input_tensor.contiguous(), target_tensor.contiguous()]

def get_init_inputs():
    """
    提供模型初始化所需的参数。
    """
    return [DELTA, REDUCTION]