hw4 bug fix (final)