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


class Model(nn.Module):
    def __init__(self, eps=1e-8):
        super(Model, self).__init__()
        self.eps = eps

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

        divergence = (p - q) ** 2 / (q + self.eps)
        return torch.sum(divergence, dim=1).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 []