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

class Model(nn.Module):
    """
    Baseline Focal Loss using PyTorch's most numerically stable API.
    """
    def __init__(self, alpha=0.25, gamma=2.0, reduction='mean'):
        super(Model, self).__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.reduction = reduction
    
    def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        # --- 关键修复1：强制转换数据类型 ---
        # 确保所有计算都在float32下进行，避免类型错误
        inputs = inputs.to(torch.float32)
        targets = targets.to(torch.float32)

        # --- 关键修复2：统一输入维度 ---
        if inputs.dim() == 2 and inputs.size(1) == 1:
            inputs = inputs.squeeze(1)

        # --- 使用PyTorch官方推荐的稳定API计算BCE ---
        # 现在inputs和targets的维度和类型都正确了
        bce = F.binary_cross_entropy_with_logits(inputs, targets, reduction='none')
        
        # 后续的Focal Loss计算保持不变
        probs = torch.sigmoid(inputs)
        pt = torch.where(targets == 1, probs, 1 - probs)
        focal_weight = self.alpha * torch.pow(1 - pt, self.gamma)
        focal_loss = focal_weight * bce
        
        if self.reduction == 'mean':
            return focal_loss.mean()
        elif self.reduction == 'sum':
            return focal_loss.sum()
        else:
            return focal_loss

batch_size = 1024
num_classes = 1

def get_inputs():
    inputs = torch.randn(batch_size, num_classes)
    targets = torch.randint(0, 2, (batch_size,))
    return [inputs, targets]

def get_init_inputs():
    return []
