forked from ccf-ai-infra/GPUCodeForces
26 lines
663 B
Python
26 lines
663 B
Python
# 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 [] |