smth:

this is what they mean:
# input is 4 x 2 x 388 x 388 raw ouputs of linear.

logits = F.log_softmax(input)

# Here, an undocumented thing is that if input is 4D, softmax is applied in dimension 1. i.e. for each pixel location, you have a softmax applied over channels.

m = nn.NLLLoss2d()

loss = m(logits, target)
# here target is of shape 4 x 1 x 388 x 388



Thank you but I don’t quite get it. Why would we use log_softmax instead of softmax? And why do we need the targets? According to the paper we only need the output feature channels of our neural network right? Later we use cross-entropy on logits and target to penalize each position from 1.
I really want to replicate the paper with a 100% authentic code and not use new techniques discovered later after the paper was released.
Following your instructions I now have:
class EnergyFunction(nn.Module):
    def __init__(self):
        super(EnergyFunction, self).__init__()

    def forward(self, logits, target):
        # logits is of shape [batch_size, feature_channels, height, width]
        logits = F.log_softmax(logits)
        m = nn.NLLLoss2d()

        # At this point:
        # logits.size() == [4, 2, 388, 388]
        # target.size() == [4, 388, 388]
        loss = m(logits, target.long())
        return loss

But i get a 1D loss as a result, this is not really what I want. My goal here is really to map the maths formulas to the code. Thank you