def METHOD_NAME(
    self,
    y_hat: Tensor,
    y: Tensor,
) -> Tensor:
    if self.model_config.return_type == ModelReturnType.raw:
        loss_fn = F.cross_entropy
    elif self.model_config.return_type == ModelReturnType.probs:
        loss_fn = F.nll_loss
        y_hat = y_hat.log()
    elif self.model_config.return_type == ModelReturnType.log_probs:
        loss_fn = F.nll_loss
    else:
        assert False
    return loss_fn(y_hat, y)