# kldivloss_torch.py
import torch
import torch.nn as nn
import torch.nn.functional as F

BATCH_SIZE = 16
DIM = 16384 * 16 

class Model(nn.Module):
   
    def forward(self, input_logits: torch.Tensor, target_prob: torch.Tensor) -> torch.Tensor:
      
        
        return F.kl_div(input_logits, target_prob, reduction='mean') 

def get_inputs():

    target_prob = torch.rand(BATCH_SIZE, DIM, dtype=torch.float32)
    target_prob = target_prob / target_prob.sum(dim=-1, keepdim=True)

    input_logits = F.log_softmax(torch.randn(BATCH_SIZE, DIM, dtype=torch.float32), dim=-1)
    
    return [input_logits, target_prob]

def get_init_inputs():
    return []