fixed bugs, added print statements, gradient check fails b/c of changing unit_num, sort of weird
Added TODOs