Hi,
if I understand correctly, the shape of label should be (batch_size)x64x128x128, because in a 2D case, the label should be (batch_size)x(height)x(width).

Yes, the highest score stands for voxel class.
There’s no problem if the output includes negative values because they aren’t applied softmax.
If you want to use cross entropy loss, you don’t need to. https://pytorch.org/docs/stable/nn.html#torch.nn.CrossEntropyLoss 1