def loss_fn(loss ,logits, labels):
    """
    Args:
        logits (torch.Tensor): Predicted logits.
        labels (torch.Tensor): Ground truth labels.
        loss (torch.nn.CrossEntropyLoss): Loss function.
    Returns:
        loss (torch.Tensor): Loss tensor.
    """
    return loss(logits.float(), labels.float())