Minor fix in assignment 1