My predictions from the UNet is having shape[1,21,224,224] but the target masks have shape [1,224,224], how should i merge prediction and use it for dice score calculation.

If i use torch.argmax() to get max prediction/labels for dice score calculation i am getting scores more then 1.

For calculating dice score i am using below function.

def dice (pred,targs):

pred=(pred>0).float()

return 2. *(pred*targs).sum()/(pred+targs).sum()