It seems the spatial sizes are permuted in the output or the target tensor.
What shapes do you expect as the height and width?
The current format would be [batch_size, channels, height, width], so the height and width are mixed in one of the mentioned tensors.