import torch


def my_multi_task_loss(y_pred, y, weights):
    """
    Multitask loss function. Computes the individual loss for each output and combine them in a weighted average
    :param y_pred: predicted output
    :param y: true labels
    :param weights: list of scalars representing the weight to apply to each loss
    :return: the weighted average loss
    """

    # mse_loss = torch.nn.MSELoss()
    l1_loss = torch.nn.L1Loss()
    xe_loss = torch.nn.CrossEntropyLoss()

    y_pred_age, y_pred_gender, y_pred_race = y_pred
    y_age, y_gender, y_race = y

    age_loss = l1_loss(y_pred_age, y_age)
    gender_loss = xe_loss(y_pred_gender, y_gender)
    race_loss = xe_loss(y_pred_race, y_race)

    avg_loss = weights[0]*age_loss + weights[1]*gender_loss + weights[2]*race_loss
    return avg_loss