Fix infinite while loop bug