Thanks for your time.
My pytorch is 1.2.0 and I find argmax always returns the last index in the above case(in CPU), so I didnot try to think about the consistency. I will follow your answer.
Again, thanks for your reply.