assignment 2 code

Minor comment update