@fmassa Does the F.log_softmax take care of the fact that we need to take softmax along the 1st axis?
For eg the input would be a tensor of size (batch_size, n_classes, H, W), we need to apply softmax along each n_classes slice for each pixel in the image sized h*w.
From what I read here 27, I’m not sure if F.log_softmax does that.
I’m currently permuting and resizing the outputs to get a slice with dimensions batch_size, h, w, n_classes and then directly using F.cross_entropy. Is this approach correct?
            # outputs.shape =(batch_size, n_classes, img_cols, img_rows) 
            outputs = outputs.permute(0, 2, 3, 1)
            # outputs.shape =(batch_size, img_cols, img_rows, n_classes) 
            outputs = outputs.resize(batch_size*img_cols*img_rows, n_classes)
            labels = labels.resize(batch_size*img_cols*img_rows)
            loss = F.cross_entropy(outputs, labels)