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

BATCH_SIZE = 4096
N_CLASSES = 1024


class Model(nn.Module):

    def __init__(self):
        super().__init__()
        self.criterion = nn.CrossEntropyLoss(reduction='mean')

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


def get_inputs():
    logits = torch.randn(BATCH_SIZE, N_CLASSES, dtype=torch.float32)
    target = torch.randint(0, N_CLASSES, (BATCH_SIZE,), dtype=torch.long)
    return [logits, target]


def get_init_inputs():
    return []