I’m writing this reply for any human in the distant future, in some not-so-distant galaxy, who needs something similar.
The module I ended up using is generic for any number of channels. It takes a (N_BATCHES, N_CHANNELS, WIDTH, DEPTH) tensor of a batch of N_BATCHES images with  N_CHANNELS each. It outputs another (N_BATCHES, N_CHANNELS, WIDTH, DEPTH) tensor containing the log probability for each batch and each channel (calculated and normalized across DEPTH and WIDTH)
class SoftmaxLogProbability2D(torch.nn.Module):
    def __init__(self):
        super(SoftmaxLogProbability2D, self).__init__()

    def forward(self, x):
        orig_shape = x.data.shape
        seq_x = []
        for channel_ix in range(orig_shape[1]):
            softmax_ = F.softmax(x[:, channel_ix, :, :].contiguous()
                                 .view((orig_shape[0], orig_shape[2] * orig_shape[3])), dim=1)\
                .view((orig_shape[0], orig_shape[2], orig_shape[3]))
            seq_x.append(softmax_.log())
        x = torch.stack(seq_x, dim=1)
        return x

It seems to work (so far).