Fix up changes from assignment 05