I have two networks 31 which both output a tensor of shape batch_size x num_classes x height x width with num_classes = 256 (there are actually just 21 classes in VOC12 but they choose background to have label 255 - I will improve on this later).
So as the label has format batch_size x 1 x height x width I can calculate cross entropy loss with:
criterion = nn.CrossEntropyLoss()

# shape: batch_size x 1 x height x width = 1 x 1 x 256 x 256
inputs = autograd.Variable(images)
# shape: batch_size x 1 x height x width = 1 x 1 x 256 x 256
targets = autograd.Variable(labels)
# shape: batch_size x 1 x height x width = 1 x 256 x 256 x 256
outputs = model(inputs)

optimizer.zero_grad()
loss = criterion(outputs.view(-1, 256), targets.view(-1))
loss.backward()
optimizer.step()

See source for context 23.
I know that outputs[0][0][i][j| corresponds to the probability that the pixel at (i, j) belongs to class 1. So if want to transform outputs of shape 1 x 256 x 256 x 256 to 1 x 1 x 256 x 256 I would need to find the maximum (probability) of every pixel and assign it to the corresponding class value.
I could do this manually by iterating over every class and pixel with numpy but I wonder if there is any better way using tensor operations?