import torch


def get_loss(input, target):
    """
    input.shape = B * C, where C is the number of class
    target.shape = B, where 0 <= target value < C
    """
    loss_fn = torch.nn.CrossEntropyLoss(reduction="none")
    return loss_fn(input, target.to(dtype=torch.int64))