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


class Model(nn.Module):
    def __init__(self, beta=0.5):
        super(Model, self).__init__()
        self.beta = beta

    def forward(self, p, q):
        p_prob = F.softmax(p, dim=1)
        q_prob = F.softmax(q, dim=1)

        term1 = p_prob.pow(self.beta)
        term2 = (self.beta - 1) * q_prob.pow(self.beta)
        term3 = self.beta * p_prob * q_prob.pow(self.beta - 1)

        sum_val = torch.sum(term1 + term2 - term3, dim=1)
        loss = sum_val / (self.beta * (self.beta - 1))

        return loss.mean()


batch_size = 32
num_classes = 1000


def get_inputs():
    p = torch.randn(batch_size, num_classes, requires_grad=True)
    q = torch.randn(batch_size, num_classes)
    return [p, q]


def get_init_inputs():
    return [0.5]