Unpacked code for better understanding:
class Softmax2DProbability(torch.nn.Module):
    def __init__(self):
        super(Softmax2DProbability, self).__init__()

    def forward(self, x):
        x_orig_shape = x.data.shape
        x_vectorised_image_size = x_orig_shape[2]*x_orig_shape[3]
        x_reshaped = x.view((x_orig_shape[0], x_vectorised_image_size))
        # now I can use torch's softmax
        soft_max_vectorised_x = F.softmax(x_reshaped, dim=1).view(x_orig_shape)
        # reshape to return
        return soft_max_vectorised_x.view(x_orig_shape)