split flattened tensor to perform class wise loss

Question

I would like to adapt my dice loss function for multi-class averaging. y_true and y_pred are one-hot 3D images (mask of label). For example, if the flattened tensors contain the