Minor tweaks before submit homework 2