Updated Assignment 2 so it works