HW4 done, fixed problem with returns computation