If you are using nn.CrossEntropyLoss and your model outputs the mentioned shape, your target should have the shape [batch_size, height, width] and contain class indices in the range [0, 3], i.e. it should not be one-hot encoded.