Assuming your input has the dimensions [height, width, channels], you would have to permute it and add a batch dimension, since PyTorch modules expect image tensors as [batch_size, channels, height, width].
The same applies to your target tensor, i.e. it should have the shape [batch_size, 21, 800, 800].
Based on your description it seems you are working on a mulit-label segemntation, i.e. each pixel position could belong to zero, one, or multiple classes.
If that’s correct, your model output should have the same shape as the target and contain logits (no non-linearity at the end). For the criterion you could use nn.BCEWithLogitsLoss.
Let me know, if that would work for you.