I am trying to calculate dice coefficient for that purpose I need to average each class. for that I need one hot encoded vector. here is code
from einops import rearrange
def dice_coef_multilabel(mask_pred,mask_gt  ):
    def compute_dice_coefficient(mask_pred,mask_gt, smooth = 0.0001):
      """Compute soerensen-dice coefficient.
      compute the soerensen-dice coefficient between the ground truth mask `mask_gt`
      and the predicted mask `mask_pred`. 

      Args:
        mask_gt: 4-dim Numpy array of type bool. The ground truth mask. [B, 1, H, W]
        mask_pred: 4-dim Numpy array of type bool. The predicted mask. [B, C, H, W]
      Returns:
        the dice coeffcient as float. If both masks are empty, the result is NaN
      """
      volume_sum = mask_gt.sum() + mask_pred.sum()
      
      volume_intersect = (mask_gt * mask_pred).sum()
      return (2*volume_intersect+smooth) / (volume_sum+smooth)
    dice=0
    n_pred_ch = mask_pred.shape[1]
    mask_pred=torch.softmax(mask_pred, 1)
    mask_gt=F.one_hot(mask_gt.long(), num_classes=n_pred_ch) #create one hot vector
    mask_gt=rearrange(mask_gt, 'd0 d1 d2 d3 d4 -> d0 (d1 d4) d2 d3 ')  #reshape one hot vector

    for ind in range(1,n_pred_ch):
        dice += compute_dice_coefficient(mask_gt[:,ind,:,:], mask_pred[:,ind,:,:])
    return dice/n_pred_ch # taking average