nn.CrossEntropyLoss expects model outputs in the shape [batch_size, nb_classes, *additional_dims], while the target should contain class indices in the range [0, nb_classes-1] and have the shape [batch_size, *additional_dims]. Assuming you are working on a multi-class segmentation use case with volumes, your output might be [batch_size, nb_classes, depth, height, width], while the target should then be [batch_size, depth, height, width]. Based on the error message I assume that the target shape might be wrong.