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


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

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

        sum_term = torch.sum((p_prob ** self.q) * (q_prob ** (1.0 - self.q)), dim=1)
        loss = (sum_term - 1.0) / (self.q - 1.0)

        return loss.mean()


batch_size = 32
num_classes = 1000


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


def get_init_inputs():
    return [0.5]