Small code fixes for Assignment 2