Assignment 2 bug fix