It depends a bit on the task you want to do.
if you want cross entropy for each pixel independently:
output_of_net = torch.rand(batch, nb_class, height, width)
label = torch.LongTensor(batch, height, width)

output_of_net_linearized = output_of_net.permute(0, 2, 3, 1).reshape(-1, output_of_net.size(1))
label_linearized = label.view(-1)
loss = criterion(output_of_net_linearized, label_linearized)