Hi there,
I’m using my last NN layer as a softmax layer for outputting a 2D normalised heatmap (probability distribution of the correct pixel in an image). I had to implement this myself for the layer
class Softmax2DProbability(torch.nn.Module):
    def __init__(self):
        super(Softmax2DProbability, self).__init__()

    def forward(self, x):
        orig_shape = x.data.shape
        return F.softmax(x.view((orig_shape[0], orig_shape[2]*orig_shape[3])), dim=1).view(orig_shape)

Assuming this layer is correct (which seems to be), how to I get a cross entropy between my NN output from this layer and the target heatmap (one hot 2D array for the pixel with the right value)?
If I reshape my tensor to use Torch’s current CrossEntropy, will autograd know automatically what to do to differentiate?
Thank you,