Solved by imaluengo in post #4 
              
              
                I personally prefer explicit code (as long as it doesn’t add a huge overhead) and in this case is pretty readable: 
data = data.view(-1).softmax(0).view(*data.shape)

It is even simpler if you don’t care about recovering the previous dimensions: 
data = data.view(-1).softmax(0)

Also, dim=-1 refers …