import torch


def dice_coefficient(pred, label, epsilon=1e-6):
    """
    Function that calculates the dice coefficient
    :param pred: N,C,H,W
    :param label: N,C,H,W
    :param epsilon:
    :return: N,C
    """
    # populate the output holder
    numerator = (pred * label).sum(dim=-1).sum(dim=-1) * 2
    denominator = pred.pow(2).sum(dim=-1).sum(dim=-1) + label.pow(2).sum(dim=-1).sum(dim=-1) + epsilon
    dice_val = numerator/denominator
    # return NaN for example where a class is not present
    no_show_indicator = label.sum(dim=-1).sum(dim=-1)==0
    return torch.where(no_show_indicator, torch.ones(dice_val.size()).to(dice_val.device) * float('nan'), dice_val)