Homework 4 small fix