Small changes in assignment 5