Let’s say you have N C H W tensors.
If you mean channel wise as in “for each pixel, a probability distribution over the channels”, then F.softmax(x, 1) is for you.
If you want “for each channel, a probability distribution over the pixels”, you should use F.softmax(x.reshape(x.size(0), x.size(1), -1), 2).view_as(x) instead. So you reshape to merge the last two dimensions, take the softmax over that, and reshape to match the input (for the latter, you know that view_as will work).
Best regards
Thomas