small code fix in assignment 6