Fix solution to exercise 2