i used torch.argmax before training the model to convert one hot encoded data ( b, class, h, w) =(16, 12, 256, 256) to (16, 256, 256). Now i can use the cross entropy loss function. Is it the right way to train a model for segmentation task?