Ah, sorry, softmax dim option cannot take tuple but only int, so you have to flatten your image before computing softmax, something like this should do the trick:
decoder_shape = decoder.shape
flatten_decoder = decoder.view(decoder_shape[0], decoder_shape[1], -1)
faltten_heat_map = torch.nn.Softmax(dim=2)(flatten_decoder )
return faltten_heat_map.view(decoder_shape)