small fix for problem 3