Dice score for prediction shape [1,21,224,224] and target shape [1,224,224]

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):
return 2. (predtargs).sum()/(pred+targs).sum()