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

BATCH_SIZE = 4096
NUM_CLASSES = 1000  # 假设分类类别数
FEATURE_DIM = NUM_CLASSES  # CrossEntropyLoss输入最后一维为类别数

class Model(nn.Module):

    def __init__(self):
        super().__init__()
        # CrossEntropyLoss 会自动包含 LogSoftmax + NLLLoss
        self.criterion = nn.CrossEntropyLoss(reduction='mean')

    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        return self.criterion(input, target)

def get_inputs():
    # CrossEntropyLoss 输入：logits [N, C]
    input_scores = torch.randn(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32)

    # 目标标签：每个样本一个类别索引（0 ~ C-1）
    target = torch.randint(0, FEATURE_DIM, (BATCH_SIZE,), dtype=torch.long)

    return [input_scores, target]

def get_init_inputs():
    return []