Fixed problem in assignment 4