Update code  fix bug