Fix code bug