Fixed some potential bugs