import torch
import torch.nn as nn

class Model(nn.Module):
    """
    PyTorch基准实现：直接从2D输入计算Dice Loss
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        """
        Applies Dice Loss to the 2D prediction and target tensors.

        Args:
            pred (torch.Tensor): Prediction tensor of shape [N, C, H, W].
            target (torch.Tensor): Target tensor of same shape as pred.

        Returns:
            torch.Tensor: Dice Loss value (scalar).
        """
        # --- 第一步：在内部展平张量 ---
        pred_flat = pred.view(-1)
        target_flat = target.view(-1)
        
        # --- 第二步：计算Dice Loss ---
        intersection = (pred_flat * target_flat).sum()
        pred_sum = pred_flat.sum()
        target_sum = target_flat.sum()
        
        epsilon = 1e-6
        dice_score = (2.0 * intersection) / (pred_sum + target_sum + epsilon)
        dice_loss = 1.0 - dice_score
        
        return dice_loss

batch_size = 32
height, width = 256, 256
channels = 1

def get_inputs():
    pred = torch.rand(batch_size, channels, height, width)
    target = torch.randint(0, 2, (batch_size, channels, height, width), dtype=torch.float32)
    return [pred, target]

def get_init_inputs():
    return []  # No special initialization inputs needed
