updated to solve assignment 2 of coursera course