import torch


def get_criterion(task_name, weights=None):
	"""
		Get the appropiate training criterion given a task
		objective

	Args:
		task_name: str, task goal [regression|binary|multilabel]
		weights: (N,) torch.Tensor of dataset labels

	Return:
		a nn.Loss loss function criterion
	"""
	if task_name == "regression":
		return torch.nn.MSELoss()

	if task_name == "multilabel":
		return torch.nn.CrossEntropyLoss(weight=weights)

	if task_name == "binary":
		return torch.nn.BCEWithLogitsLoss()

	raise IOError("Undefined task - Missing criterion!")