import torch


def weighted_nll_loss(inpt: torch.Tensor, target: torch.Tensor, weight: torch.Tensor):
    """What `torch.nll_loss` with `reduction='weighted_average'` should be"""
    return (
        torch.nn.functional.nll_loss(inpt, target, weight=weight, reduction="sum")
        / weight.take(target).sum()
    )