problem 5: fix bug