Make changes according to the comments from the review for Assignment 2