Thanks for the detailed answer @ClementPinard. Your answer is clearly solving the problem I posted however is does not incorporate well with my original goal which is implementing a dice loss. I need to filter out invalid label, it’s the reason I need to apply a mask on the one-hot encoded labels, but I also need to filter the output tensor. However I don’t think I can do the same array modification and keep autograd working correctly. Here is the original code
def dice_loss3(output, target, weights=1, ignore_index=None):
    output = output.exp()
    encoded_target = output.data.clone().zero_()
    if ignore_index is not None:
        # mask of invalid label
        mask = target == ignore_index
        # clone target to not affect the variable ?
        filtered_target = target.clone()
        # replace invalid label with whatever legal index value
        filtered_target[mask] = 0
        # one hot encoding
        encoded_target.scatter_(1, filtered_target.unsqueeze(1), 1)
        # expand the mask for the encoded target array
        mask = mask.unsqueeze(1).expand(output.data.size())
    else:
        encoded_target.scatter_(1, target.unsqueeze(1), 1)
    encoded_target = Variable(encoded_target)

    assert output.size() == encoded_target.size(), "Input sizes must be equal."
    assert output.dim() == 4, "Input must be a 4D Tensor."

    o = output[~mask].view(-1, output.size(1))
    t = encoded_target[~mask].view(-1, output.size(1))
    numerator = o * t
    denominator = o.pow(2) + t

    dice = 2 * (numerator / denominator) * weights
    return dice.sum() 

I tried to adapt :
def dice_loss3(output, target, weights=1, ignore_index=None):
    b, c, h, w = output.size()
    output = output.exp()
    output = output.permute(1, 0, 2, 3)
    encoded_target = output.data.new(c, b, h, w).zero_()
    if ignore_index is not None:
        # mask of invalid label
        mask = target == ignore_index
        # clone target to not affect the variable ?
        filtered_target = target.clone()
        # replace invalid label with whatever legal index value
        filtered_target[mask] = 0
        # one hot encoding
        encoded_target.scatter_(0, filtered_target.unsqueeze(0), 1)
        # expand the mask for the encoded target array
        mask = mask.unsqueeze(0).expand(encoded_target.size())
    else:
        encoded_target.scatter_(1, target.unsqueeze(1), 1)
    encoded_target = Variable(encoded_target)

    assert output.size() == encoded_target.size(), "Input sizes must be equal."
    assert output.dim() == 4, "Input must be a 4D Tensor."

    o = output[~mask].view(output.size(0), -1)
    t = encoded_target[~mask].view(output.size(0), -1)
    numerator = o * t
    denominator = o.pow(2) + t

    dice = 2 * (numerator.sum(1) / denominator.sum(1)) * weights
    return dice.sum()

This is clearly wrong as I am having dice values over 5000, it shouldn’t be higher than 1. The formula is simple (output*one_hot_targets) / (output.pow(2) + one_hot_targets)
If you have any idea what could be wrong, it would be really nice, I tried so many hacks to solve this… Thanks