Solved by ptrblck in post #2 
              
              
                Could you print the shape of the tensor you are passing to F.softmax? 
It seems dim2 is just missing. Are you adding the batch dimension to your data in the test case? This is often forgotten and yields these kind of errors. 
You can call functions like softmax on “negative” dimensions to use revers…