The error is raised, since you are trying to use a scalar target (single class index for each sample) on an image-like model output, as you’ve already explained.
This won’t work out of the box and I don’t know, how TensorFlow manages this use case without a hidden reduction somewhere.
If you are working on a multi-class classification, which seems to be the case based on your target, you would have to provide class logits as the model output in the shape [batch_size, nb_classes], where each value in dim1 represents the logit for the corresponding class index.
Usually this is done by flattening the conv layer output and by passing it through a final classifier e.g. a linear layer.
If you don’t want to use a linear layer, you would have to make sure your last activation output from lastcnn returns a single pixel output in the shape [batch_size, nb_classes, 1, 1].
If your output already has this shape, you could use squeeze to provide the desired output via:
output = output.squeeze(3).squeeze(2)

and pass it to the criterion.
If this activation has more than a single pixel, you could either change the model architecture by increasing the stride etc. or add an adaptive pooling layer, which would output the 1x1 activation.