import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self, reduction='mean'):
        super().__init__()
        self.reduction = reduction
        self.mse_loss = nn.MSELoss(reduction=reduction)

    def forward(self, logits_student: torch.Tensor, logits_teacher: torch.Tensor) -> torch.Tensor:
        return self.mse_loss(logits_student, logits_teacher)


batch_size = 256
num_classes = 1000


def get_inputs():
    logits_student = torch.randn(batch_size, num_classes, dtype=torch.float32)
    logits_teacher = torch.randn(batch_size, num_classes, dtype=torch.float32)
    return [logits_student, logits_teacher]


def get_init_inputs():
    return ['mean']