Use of dim=0/1 in pytorch and nn.softmax?
When using nn.softmax(), we use dim=1 or 0. Here dim=0 should mean row according to intuition but seems it means along the column. Is this true?
&gt;&gt;&gt; x = torch.tensor([[1,2],[3,4]],dtype=torch.float)
&gt;&gt;&gt; F.softmax(x,dim=0)
tensor([[0.1192, 0.1192],
        [0.8808, 0.8808]])
&gt;&gt;&gt; F.softmax(x,dim=1)
tensor([[0.2689, 0.7311],
        [0.2689, 0.7311]])

Here when dim=0, probabilities along the columns sum to 1. Similarly when dim=1 probabilities along the rows sum to 1. Can someone explain how dim is used in PyTorch?

Indeed, in the 2D case: row refers to axis=0, while column refers to axis=1.
The dim option specifies along which dimension the softmax is apply, i.e. summing back on that same axis will lead to 1s:
&gt;&gt;&gt; x = torch.arange(1, 7, dtype=float).reshape(2,3)
tensor([[1., 2., 3.],
        [4., 5., 6.]], dtype=torch.float64)

On axis=0:
&gt;&gt;&gt; F.softmax(x, dim=0).sum(0)
tensor([1.0000, 1.0000, 1.0000], dtype=torch.float64)

On axis=1:
&gt;&gt;&gt; F.softmax(x, dim=1).sum(1)
&gt;&gt;&gt; tensor([1.0000, 1.0000], dtype=torch.float64)

This is the expected behavior for torch.nn.functional.softmax

[...]
Parameters:

dim (int) – A dimension along which Softmax will be computed (so every slice along dim will sum to 1).


