Hi @edshkim98,
I think it could be because of the diceCorrect_g = (pred * label)[:,0], because pred is of shape [B*N, 1] and represents the models output for the selected class i, while label is still of shape [B*N, C]. So multiplying each other you get diceCorrect_g, which should be of shape [B*N, C] and then you select the first element of dim 1, which probably is mostly zero.
Either you set label = label_g[:, i] (where i denotes your class) or I think you can actually remove the for loop totally and just do diceCorrect_g = (label_g * softmax(prediction_g, dim=-1)).sum() and dicePrediction_g = dicePrediction_g .sum() diceLabel_g = diceLabel_g .sum()