Update code follow comment from Mr SonTH.