argmax wrapper fix, please document argmax, its outputs are non-intuitive