Final Assignment 4 tweaks