import torch


def _loss_fn(logits, labels, weights):
    """
    Cross-entropy loss
    """
    loss = torch.nn.CrossEntropyLoss(reduction='none')(logits, labels)
    return torch.sum(loss * weights)