Ok, I had the time to have a quick look. I hope I’m not wrong, but from what I understood the softmax layer is used per map, and not per pixel. If I understood right looking at the generated confidence-maps page 6, the softmax is applied on the spatial dimension, not across the channels, which mean you want a spatial distribution and not a class distribution.
So If you follow the pytorch convention for the dmension of your images (Batch, channels, x, y), then you should apply your softmax as:
return torch.nn.Softmax(dim=(2,3))(decoder)
I hope it should be right like that.