Pytorch softmax: What dimension to use? : <p>The function <code>torch.nn.functional.softmax</code> takes two parameters: <code>input</code> and <code>dim</code>.  According to its documentation, the softmax operation is applied to all slices of <code>input</code> along the specified <code>dim</code>, and will rescale them so that the elements lie in the range <code>(0, 1)</code> and sum to 1. </p>

<p>Let input be:</p>

<pre><code>input = torch.randn((3, 4, 5, 6))
</code></pre>

<p>Suppose I want the following, so that every entry in that array is 1:</p>

<pre><code>sum = torch.sum(input, dim = 3) # sum's size is (3, 4, 5, 1)
</code></pre>

<p>How should I apply softmax?</p>

<pre><code>softmax(input, dim = 0) # Way Number 0
softmax(input, dim = 1) # Way Number 1
softmax(input, dim = 2) # Way Number 2
softmax(input, dim = 3) # Way Number 3
</code></pre>

<p>My intuition tells me that is the last one, but I am not sure. English is not my first language and the use of the word <code>along</code> seemed confusing to me because of that.</p>

<p>I am not very clear on what "along" means, so I will use an example that could clarify things.  Suppose we have a tensor of size (s1, s2, s3, s4), and I want this to happen</p>
