import torch


def get_criterion(config):
    """
    Creates the torch criterion for optimization
    """
    if config["name"] == "cross_entropy":
        return torch.nn.CrossEntropyLoss()
    else:
        raise NotImplementedError